Switch to manual for loop to support select statements.

docs/collections_doc.md

@@ -81,7 +81,9 @@ A new list with all unique elements from `iterable`.
 collections.flatten(<a href="#collections.flatten-iterable">iterable</a>)
 </pre>
 
-Collapses the first dimension of the iterable and returns the result as a list.
+Flattens an iterable to the sum of those iterables.
+
+This behaves like collapsing the first dimension of the iterable.
 
 **PARAMETERS**
 

lib/collections.bzl

@@ -66,7 +66,9 @@ def _uniq(iterable):
     return list(unique_elements.keys())
 
 def _flatten(iterable):
-    """Collapses the first dimension of the iterable and returns the result as a list.
+    """Flattens an iterable to the sum of those iterables.
+
+    This behaves like collapsing the first dimension of the iterable.
 
     Args:
       iterable: An iterable of iterables to be collapsed to a list.
@@ -74,11 +76,10 @@ def _flatten(iterable):
     Returns:
       A new list with the collapsed elements from `iterable`.
     """
-    return [
-        element
-        for subiterable in iterable
-        for element in subiterable
-    ]
+    result = []
+    for element in iterable:
+      result += element
+    return result
 
 collections = struct(
     after_each = _after_each,

tests/collections_tests.bzl

@@ -105,9 +105,16 @@ uniq_test = unittest.make(_uniq_test)
 
 def _flatten_test(ctx):
     env = unittest.begin(ctx)
-    asserts.equals(env, collections.flatten([[0, 1], [2, 3]]), [0, 1, 2, 3])
-    asserts.equals(env, collections.flatten([]), [])
-    asserts.equals(env, collections.flatten([[], [0, 1]]), [0, 1])
+    asserts.equals(env, [0, 1, 2, 3], collections.flatten([[0, 1], [2, 3]]))
+    asserts.equals(env, [], collections.flatten([]), [])
+    asserts.equals(env, [0, 1], collections.flatten([[], [0, 1]]))
+    asserts.equals(env,
+        '[] + select({"//conditions:default": [0, 1]}) + select({"//conditions:default": [2, 3]})',
+        str(collections.flatten([select({
+            "//conditions:default": [0, 1],
+        }), select({
+            "//conditions:default": [2, 3],
+        })])))
 
     return unittest.end(env)