Support tuple destructuring in implied for's. (#220)

* Add docs and coverage for tuple-target comprehension lowering
# Conflicts:
#	README.md
#	docs/source/generic/query_structure.md
#	func_adl/ast/syntatic_sugar.py

* Fix up

Author: gordonwatts

README.md

@@ -58,6 +58,7 @@ There are several python expressions and idioms that are translated behind your
 |List Comprehension | `[j.pt() for j in jets]` | `jets.Select(lambda j: j.pt())` |
 |List Comprehension | `[j.pt() for j in jets if abs(j.eta()) < 2.4]` | `jets.Where(lambda j: abs(j.eta()) < 2.4).Select(lambda j: j.pt())` |
 |Multi-generator comprehension|`[j.pt() + e.pt() for j in jets for e in electrons]`|`jets.SelectMany(lambda j: electrons.Select(lambda e: j.pt() + e.pt()))`|
+|Destructuring List Comprehension|`[a+b for a,b in pairs]`|`pairs.Select(lambda tmp: tmp[0] + tmp[1])`|
 |Literal List Comprehension|`[i for i in [1, 2, 3]]`|`[1, 2, 3]`|
 | Data Classes<br>(typed) | `@dataclass`<br>`class my_data:`<br>`x: ObjectStream[Jets]`<br><br>`Select(lambda e: my_data(x=e.Jets()).x)` | `Select(lambda e: {'x': e.Jets()}.x)` |
 | Named Tuple<br>(typed) | `class my_data(NamedTuple):`<br>`x: ObjectStream[Jets]`<br><br>`Select(lambda e: my_data(x=e.Jets()).x)` | `Select(lambda e: {'x': e.Jets()}.x)` |

docs/source/generic/query_structure.md

@@ -42,6 +42,9 @@ expressions:
 - List/generator comprehensions over streams are lowered to query operators.
   Multi-generator forms (`for ... for ...`) flatten outer levels via `.SelectMany(...)`
   so the stream shape matches Python iteration semantics.
+- Comprehensions with tuple/list destructuring targets are supported. For non-literal
+  iterables, destructured names are rewritten as indexed access on an internal temporary
+  value (for example, `[a+b for a,b in pairs]` behaves like `pairs.Select(lambda tmp: tmp[0] + tmp[1])`).
 - List comprehensions over literal iterables are expanded directly. For example,
   `[i for i in [1, 2, 3]]` becomes `[1, 2, 3]`.
 - `any`/`all` over literal lists/tuples are reduced to boolean `or`/`and` expressions.

func_adl/ast/syntatic_sugar.py

@@ -70,6 +70,37 @@ def _target_bindings(
                 f" - {ast.unparse(node)}"
             )
 
+        def _target_bindings_from_value(
+            self, target: ast.AST, value: ast.expr, node: ast.AST
+        ) -> Dict[str, ast.expr]:
+            """Build bindings for comprehension targets, using indexing for non-literals."""
+
+            # Reuse literal destructuring where possible to preserve exact tuple/list nodes.
+            literal_bindings = self._target_bindings(target, value, node)
+            if literal_bindings is not None:
+                return literal_bindings
+
+            if isinstance(target, ast.Name):
+                return {target.id: copy.deepcopy(value)}
+
+            if isinstance(target, (ast.Tuple, ast.List)):
+                bindings: Dict[str, ast.expr] = {}
+                for index, target_elt in enumerate(target.elts):
+                    element_value = ast.Subscript(
+                        value=copy.deepcopy(value),
+                        slice=ast.Constant(value=index),
+                        ctx=ast.Load(),
+                    )
+                    bindings.update(
+                        self._target_bindings_from_value(target_elt, element_value, node)
+                    )
+                return bindings
+
+            raise ValueError(
+                f"Comprehension variable must be a name or tuple/list, but found {target}"
+                f" - {ast.unparse(node)}"
+            )
+
         def _substitute_names(self, expr: ast.expr, bindings: Dict[str, ast.expr]) -> ast.expr:
             class _name_replacer(ast.NodeTransformer):
                 def __init__(self, loop_bindings: Dict[str, ast.expr]):
@@ -371,25 +402,42 @@ def resolve_generator(
             """
             a = node
             generator_count: int = len(generators)
+            temp_counter = 0
             for index, c in enumerate(reversed(generators)):
                 target = c.target
-                if not isinstance(target, ast.Name):
-                    # Keep original comprehension for unsupported lowering cases.
-                    return node
                 if c.is_async:
                     raise ValueError(f"Comprehension can't be async - {ast.unparse(node)}.")
                 source_collection = c.iter
 
+                if isinstance(target, ast.Name):
+                    lambda_arg_name = target.id
+                elif isinstance(target, (ast.Tuple, ast.List)):
+                    lambda_arg_name = f"__fa_tmp_{temp_counter}"
+                    temp_counter += 1
+                else:
+                    raise ValueError(
+                        "Comprehension variable must be a name or tuple/list, "
+                        f"but found {target} - {ast.unparse(node)}"
+                    )
+
+                target_bindings = self._target_bindings_from_value(
+                    target,
+                    ast.Name(id=lambda_arg_name, ctx=ast.Load()),
+                    node,
+                )
+
                 # Turn the if clauses into Where statements
                 for a_if in c.ifs:
-                    where_function = lambda_build(target.id, a_if)
+                    where_body = self._substitute_names(a_if, target_bindings)
+                    where_function = lambda_build(lambda_arg_name, where_body)
                     source_collection = ast.Call(
                         func=ast.Attribute(attr="Where", value=source_collection, ctx=ast.Load()),
                         args=[where_function],
                         keywords=[],
                     )
 
-                lambda_function = lambda_build(target.id, lambda_body)
+                rewritten_lambda_body = self._substitute_names(lambda_body, target_bindings)
+                lambda_function = lambda_build(lambda_arg_name, rewritten_lambda_body)
                 use_select_many = generator_count > 1 and index > 0
                 a = ast.Call(
                     func=ast.Attribute(

tests/ast/test_syntatic_sugar.py

@@ -90,13 +90,38 @@ def test_resolve_3generator_list_comp_flattened_shape():
     ) == ast.dump(a_new)
 
 
-def test_resolve_bad_iterator():
-    a = ast.parse("[j.pt() for idx,j in enumerate(jets)]")
+def test_resolve_tuple_target():
+    a = ast.parse("[a+b for a,b in pairs]")
     a_new = resolve_syntatic_sugar(a)
 
-    # Unsupported lowering (tuple target with non-literal source) should be
-    # preserved for downstream processing.
-    assert ast.unparse(a_new) == ast.unparse(a)
+    assert ast.dump(
+        ast.parse("pairs.Select(lambda __fa_tmp_0: __fa_tmp_0[0] + __fa_tmp_0[1])")
+    ) == ast.dump(a_new)
+
+
+def test_resolve_tuple_target_nested_with_if():
+    a = ast.parse("[a + c for (a, (b, c)) in triples if b > 0 if c < 10]")
+    a_new = resolve_syntatic_sugar(a)
+
+    assert ast.dump(
+        ast.parse(
+            "triples.Where(lambda __fa_tmp_0: __fa_tmp_0[1][0] > 0)"
+            ".Where(lambda __fa_tmp_0: __fa_tmp_0[1][1] < 10)"
+            ".Select(lambda __fa_tmp_0: __fa_tmp_0[0] + __fa_tmp_0[1][1])"
+        )
+    ) == ast.dump(a_new)
+
+
+def test_resolve_tuple_target_from_enumerate_with_if():
+    a = ast.parse("[idx + j.pt() for idx, j in enumerate(jets) if idx > 0]")
+    a_new = resolve_syntatic_sugar(a)
+
+    assert ast.dump(
+        ast.parse(
+            "enumerate(jets).Where(lambda __fa_tmp_0: __fa_tmp_0[0] > 0)"
+            ".Select(lambda __fa_tmp_0: __fa_tmp_0[0] + __fa_tmp_0[1].pt())"
+        )
+    ) == ast.dump(a_new)
 
 
 def test_resolve_no_async():