Support dict unpacking with conditional expression (#199)

* Resolve conditional dict expansions when test constant

* Fix up type and flake8 errors

* Fix up flake8 tests

* Update agents for better testing

* Improve conditional dict error msg and add nested ifexp test (#200)

* Remove unneeded recursion

Author: gordonwatts

AGENTS.md

@@ -2,3 +2,5 @@
 
 * Prefer using `Optional[int]` rather than `int | None`.
 * If there is an error the user will see (e.g. `ValueError`), make sure there is enough context in the message for the user. For example, if it is during an expression parse, include the `ast.unparse(a)` in the error message.
+* Before finishing, make sure `flake8` runs without errors on source and test files.
+* Before finishing, also make sure `black` runs without modifying files.

func_adl/ast/syntatic_sugar.py

@@ -173,6 +173,60 @@ def convert_call_to_dict(
                 values=arg_values,
             )
 
+        def _merge_into(self, target: ast.expr, add: ast.Dict) -> ast.expr:
+            """Merge ``add`` dictionary into ``target`` which may itself be a
+            dictionary or an if-expression containing dictionaries."""
+
+            if isinstance(target, ast.Dict):
+                return ast.Dict(keys=target.keys + add.keys, values=target.values + add.values)
+            else:
+                return target
+
+        def visit_Dict(self, node: ast.Dict) -> Any:
+            """Flatten ``**`` expansions in dictionary literals.
+
+            If the starred expression is a dictionary it is merged directly.  If
+            the expression is an ``if`` with both branches being dictionaries and
+            the test is a constant, it is resolved at transformation time.  If
+            the test is not resolvable, an error is raised as the back end cannot
+            translate it."""
+
+            a = self.generic_visit(node)
+            assert isinstance(a, ast.Dict)
+
+            base_keys: List[Optional[ast.expr]] = []
+            base_values: List[ast.expr] = []
+            expansions: List[ast.expr] = []
+            for k, v in zip(a.keys, a.values):
+                if k is None:
+                    expansions.append(v)
+                else:
+                    base_keys.append(k)
+                    base_values.append(v)
+
+            result: ast.AST = ast.Dict(keys=base_keys, values=base_values)
+
+            for e in expansions:
+                if isinstance(e, ast.Dict):
+                    result = self._merge_into(result, e)
+                elif (
+                    isinstance(e, ast.IfExp)
+                    and isinstance(e.body, ast.Dict)
+                    and isinstance(e.orelse, ast.Dict)
+                ):
+                    if isinstance(e.test, ast.Constant):
+                        branch = e.body if bool(e.test.value) else e.orelse
+                        result = self._merge_into(result, branch)
+                    else:
+                        raise ValueError(
+                            "Conditional dictionary expansion requires a constant test"
+                            f" - {ast.unparse(e)}"
+                        )
+                else:
+                    return a
+
+            return result
+
         def visit_Call(self, node: ast.Call) -> Any:
             """
             This method checks if the call is to a dataclass or a named tuple and converts

tests/ast/test_meta_data.py

@@ -89,13 +89,9 @@ def test_query_metadata_burried():
 
 
 def test_query_metadata_updated():
-    '''This is testing code in QMetaData, but we need lookup_query_metadata which we are
-    testing in this file'''
-    r = (
-        my_event()
-        .QMetaData({"one": "two"})
-        .QMetaData({"one": "three"})
-    )
+    """This is testing code in QMetaData, but we need lookup_query_metadata which we are
+    testing in this file"""
+    r = my_event().QMetaData({"one": "two"}).QMetaData({"one": "three"})
 
     assert lookup_query_metadata(r, "one") == "three"
 

tests/ast/test_syntatic_sugar.py

@@ -282,3 +282,78 @@ def test_resolve_compare_list_wrong_order():
     a = ast.parse("[31, 51] in p.absPdgId()")
     with pytest.raises(ValueError, match="Right side"):
         resolve_syntatic_sugar(a)
+
+
+def test_resolve_dict_star_merge():
+    """Dictionary unpacking should be flattened"""
+
+    a = ast.parse("{'n': e.EventNumber(), **{'m': e.EventNumber()}}").body[0].value  # type: ignore
+    a_resolved = resolve_syntatic_sugar(a)
+
+    expected = (
+        ast.parse("{'n': e.EventNumber(), 'm': e.EventNumber()}").body[0].value  # type: ignore
+    )
+    assert ast.unparse(a_resolved) == ast.unparse(expected)
+
+
+def test_resolve_dict_star_ifexp_true():
+    """Conditional dictionary unpacking should resolve when condition is True"""
+
+    a = (
+        ast.parse("{'n': e.EventNumber(), **({'m': e.EventNumber()} if True else {})}")
+        .body[0]
+        .value  # type: ignore
+    )
+    a_resolved = resolve_syntatic_sugar(a)
+
+    expected = (
+        ast.parse("{'n': e.EventNumber(), 'm': e.EventNumber()}").body[0].value  # type: ignore
+    )
+    assert ast.unparse(a_resolved) == ast.unparse(expected)
+
+
+def test_resolve_dict_star_ifexp_false():
+    """Conditional dictionary unpacking should resolve when condition is False"""
+
+    a = (
+        ast.parse("{'n': e.EventNumber(), **({'m': e.EventNumber()} if False else {})}")
+        .body[0]
+        .value  # type: ignore
+    )
+    a_resolved = resolve_syntatic_sugar(a)
+
+    expected = ast.parse("{'n': e.EventNumber()}").body[0].value  # type: ignore
+    assert ast.unparse(a_resolved) == ast.unparse(expected)
+
+
+def test_resolve_dict_star_ifexp_unknown():
+    """Unresolvable conditions should result in an error"""
+
+    a = (
+        ast.parse("{'n': e.EventNumber(), **({'m': e.EventNumber()} if cond else {})}")
+        .body[0]
+        .value  # type: ignore
+    )
+    with pytest.raises(ValueError, match="Conditional dictionary"):
+        resolve_syntatic_sugar(a)
+
+
+def test_resolve_dict_star_ifexp_nested():
+    """Nested conditional dictionary unpacking should resolve correctly"""
+
+    a = (
+        ast.parse(
+            "{'n': e.EventNumber(), **({'m': e.EventNumber(), **({'o': e.EventNumber()} "
+            "if True else {})} if True else {})}"
+        )
+        .body[0]
+        .value  # type: ignore
+    )
+    a_resolved = resolve_syntatic_sugar(a)
+
+    expected = (
+        ast.parse("{'n': e.EventNumber(), 'm': e.EventNumber(), 'o': e.EventNumber()}")
+        .body[0]
+        .value  # type: ignore
+    )
+    assert ast.unparse(a_resolved) == ast.unparse(expected)

tests/test_util_types.py

@@ -314,7 +314,7 @@ def test_get_method_and_class_inherrited_template():
 
     class bogus_1(Generic[T]):
         def fork(self) -> T:
-            ...
+            pass
 
     class bogus_2(bogus_1[int]):
         pass
@@ -325,7 +325,7 @@ class bogus_2(bogus_1[int]):
 def test_get_method_and_class_iterable():
     class bogus:
         def fork(self):
-            ...
+            pass
 
     assert get_method_and_class(Iterable[bogus], "fork") is None