Repository: iris-hep/func_adl

Issue:
Named Tuples and data classes
Data classes are very useful as typed objects to be passed around.

As a way to avoid problems with dictionaries, allow typed data classes (or if untyped, then untyped!).

This is being done in part to help with the [idap test](https://github.com/iris-hep/idap-200gbps-atlas/issues/47).

Discussion:

Patch:
diff --git a/func_adl/ast/syntatic_sugar.py b/func_adl/ast/syntatic_sugar.py
index fc2658d..1e788db 100644
--- a/func_adl/ast/syntatic_sugar.py
+++ b/func_adl/ast/syntatic_sugar.py
@@ -1,4 +1,6 @@
 import ast
+import inspect
+from dataclasses import is_dataclass
 from typing import Any, List
 
 from func_adl.util_ast import lambda_build
@@ -9,6 +11,7 @@ def resolve_syntatic_sugar(a: ast.AST) -> ast.AST:
 
     * List comprehensions are turned into `Select` statements
     * Generator expressions are turned into `Select` statements
+    * A data class is converted into a dictionary.
 
     Args:
         a (ast.AST): The AST to scan for syntatic sugar
@@ -83,4 +86,72 @@ def resolve_syntatic_sugar(a: ast.AST) -> ast.AST:
 
             return a
 
+        def convert_call_to_dict(
+            self, a: ast.Call, node: ast.AST, sig_arg_names: List[str]
+        ) -> ast.AST:
+            """Translate a data class into a dictionary.
+
+            Args:
+                a (ast.Call): The call node representing the data class instantiation
+                node (ast.AST): The original AST node
+
+            Returns:
+                ast.AST: The reformed AST as a dictionary
+            """
+            if len(sig_arg_names) < (len(a.args) + len(a.keywords)):
+                assert isinstance(a.func, ast.Constant)
+                raise ValueError(
+                    f"Too many arguments for dataclass {a.func.value} - {ast.unparse(node)}."
+                )
+
+            arg_values = a.args
+            arg_names = [ast.Constant(value=n) for n in sig_arg_names[: len(arg_values)]]
+            arg_lookup = {a.arg: a.value for a in a.keywords}
+            for name in sig_arg_names[len(arg_values) :]:
+                if name in arg_lookup:
+                    arg_values.append(arg_lookup[name])
+                    arg_names.append(ast.Constant(value=name))
+
+            for name in arg_lookup.keys():
+                if name not in sig_arg_names:
+                    assert isinstance(a.func, ast.Constant)
+                    raise ValueError(
+                        f"Argument {name} not found in dataclass {a.func.value}"
+                        f" - {ast.unparse(node)}."
+                    )
+
+            return ast.Dict(
+                keys=arg_names,  # type: ignore
+                values=arg_values,
+            )
+
+        def visit_Call(self, node: ast.Call) -> Any:
+            """
+            This method checks if the call is to a dataclass or a named tuple and converts
+            the call to a dictionary representation if so.
+
+            Args:
+                node (ast.Call): The AST Call node to visit.
+            Returns:
+                Any: The transformed node if it matches the criteria, otherwise the original node.
+            """
+            a = self.generic_visit(node)
+
+            if isinstance(a, ast.Call) and isinstance(a.func, ast.Constant):
+                if is_dataclass(a.func.value):
+                    assert isinstance(a.func, ast.Constant)
+
+                    # We have a dataclass. Turn it into a dictionary
+                    signature = inspect.signature(a.func.value)  # type: ignore
+                    sig_arg_names = [p.name for p in signature.parameters.values()]
+
+                    return self.convert_call_to_dict(a, node, sig_arg_names)
+
+                elif hasattr(a.func.value, "_fields"):
+                    # We have a named tuple. Turn it into a dictionary
+                    arg_names = [n for n in a.func.value._fields]
+
+                    return self.convert_call_to_dict(a, node, arg_names)
+            return a
+
     return syntax_transformer().visit(a)
diff --git a/func_adl/util_ast.py b/func_adl/util_ast.py
index 5fac886..d4af1ba 100644
--- a/func_adl/util_ast.py
+++ b/func_adl/util_ast.py
@@ -4,6 +4,7 @@ import ast
 import inspect
 import tokenize
 from collections import defaultdict
+from dataclasses import is_dataclass
 from enum import Enum
 from types import ModuleType
 from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union, cast
@@ -305,10 +306,13 @@ class _rewrite_captured_vars(ast.NodeTransformer):
         return v
 
     def visit_Call(self, node: ast.Call) -> Any:
-        "If the rewritten call turns into an actual function, then we have to bail"
+        "If the rewritten call turns into an actual function, then we have to bail,"
         old_func = node.func
         rewritten_call = cast(ast.Call, super().generic_visit(node))
-        if isinstance(rewritten_call.func, ast.Constant):
+        if isinstance(rewritten_call.func, ast.Constant) and not (
+            is_dataclass(rewritten_call.func.value)
+            or hasattr(rewritten_call.func.value, "_fields")
+        ):
             rewritten_call.func = old_func
 
         return rewritten_call


Test patch:
diff --git a/tests/ast/test_syntatic_sugar.py b/tests/ast/test_syntatic_sugar.py
index 75a85ac..d425102 100644
--- a/tests/ast/test_syntatic_sugar.py
+++ b/tests/ast/test_syntatic_sugar.py
@@ -1,8 +1,12 @@
 import ast
+from collections import namedtuple
+from dataclasses import dataclass
 
 import pytest
 
+from func_adl import ObjectStream
 from func_adl.ast.syntatic_sugar import resolve_syntatic_sugar
+from func_adl.util_ast import parse_as_ast
 
 
 def test_resolve_normal_expression():
@@ -72,3 +76,181 @@ def test_resolve_no_async():
         resolve_syntatic_sugar(a)
 
     assert "can't be async" in str(e)
+
+
+class LocalJet:
+    def pt(self) -> float:
+        return 0.0
+
+
+def test_resolve_dataclass_no_args():
+    "Make sure a dataclass becomes a dictionary"
+
+    @dataclass
+    class dc_1:
+        x: ObjectStream[LocalJet]
+        y: ObjectStream[LocalJet]
+
+    a = parse_as_ast(lambda e: dc_1(e.Jets(), e.Electrons()).x.Select(lambda j: j.pt())).body
+    a_resolved = resolve_syntatic_sugar(a)
+
+    a_expected = ast.parse("{'x': e.Jets(), 'y': e.Electrons()}.x.Select(lambda j: j.pt())")
+    assert ast.unparse(a_resolved) == ast.unparse(a_expected)
+
+
+def test_resolve_dataclass_named_args():
+    "Make sure a dataclass becomes a dictionary"
+
+    @dataclass
+    class dc_1:
+        x: ObjectStream[LocalJet]
+        y: ObjectStream[LocalJet]
+
+    a = parse_as_ast(lambda e: dc_1(x=e.Jets(), y=e.Electrons()).x.Select(lambda j: j.pt())).body
+    a_resolved = resolve_syntatic_sugar(a)
+
+    a_expected = ast.parse("{'x': e.Jets(), 'y': e.Electrons()}.x.Select(lambda j: j.pt())")
+    assert ast.unparse(a_resolved) == ast.unparse(a_expected)
+
+
+def test_resolve_dataclass_mixed_args():
+    "Make sure a dataclass becomes a dictionary"
+
+    @dataclass
+    class dc_1:
+        x: ObjectStream[LocalJet]
+        y: ObjectStream[LocalJet]
+
+    a = parse_as_ast(lambda e: dc_1(e.Jets(), y=e.Electrons()).x.Select(lambda j: j.pt())).body
+    a_resolved = resolve_syntatic_sugar(a)
+
+    a_expected = ast.parse("{'x': e.Jets(), 'y': e.Electrons()}.x.Select(lambda j: j.pt())")
+    assert ast.unparse(a_resolved) == ast.unparse(a_expected)
+
+
+def test_resolve_dataclass_too_few_args():
+    "Make sure a dataclass becomes a dictionary"
+
+    @dataclass
+    class dc_1:
+        x: ObjectStream[LocalJet]
+        y: ObjectStream[LocalJet]
+
+    a = parse_as_ast(lambda e: dc_1(e.Jets()).x.Select(lambda j: j.pt())).body  # type: ignore
+    a_resolved = resolve_syntatic_sugar(a)
+
+    a_expected = ast.parse("{'x': e.Jets()}.x.Select(lambda j: j.pt())")
+    assert ast.unparse(a_resolved) == ast.unparse(a_expected)
+
+
+def test_resolve_dataclass_too_many_args():
+    "Make sure a dataclass becomes a dictionary"
+
+    @dataclass
+    class dc_1:
+        x: ObjectStream[LocalJet]
+        y: ObjectStream[LocalJet]
+
+    with pytest.raises(ValueError):
+        a = parse_as_ast(
+            lambda e: dc_1(e.Jets(), e.Electrons(), e.Muons(), e.Jets()).x.Select(  # type: ignore
+                lambda j: j.pt()
+            )  # type: ignore
+        ).body
+        resolve_syntatic_sugar(a)
+
+
+def test_resolve_dataclass_bad_args():
+    "Make sure a dataclass becomes a dictionary"
+
+    @dataclass
+    class dc_1:
+        x: ObjectStream[LocalJet]
+        y: ObjectStream[LocalJet]
+
+    with pytest.raises(ValueError):
+        a = parse_as_ast(
+            lambda e: dc_1(z=e.Jets()).x.Select(lambda j: j.pt())  # type: ignore
+        ).body
+        resolve_syntatic_sugar(a)
+
+
+def test_resolve_named_tuple_no_args():
+    "Make sure a named tuple becomes a dictionary"
+
+    nt_1 = namedtuple("nt_1", ["x", "y"])
+
+    a = parse_as_ast(lambda e: nt_1(e.Jets(), e.Electrons()).x.Select(lambda j: j.pt())).body
+    a_resolved = resolve_syntatic_sugar(a)
+
+    a_expected = ast.parse("{'x': e.Jets(), 'y': e.Electrons()}.x.Select(lambda j: j.pt())")
+    assert ast.unparse(a_resolved) == ast.unparse(a_expected)
+
+
+def test_resolve_named_tuple_typed():
+    "Make sure a named tuple becomes a dictionary"
+
+    from typing import NamedTuple
+
+    class nt_1(NamedTuple):
+        x: ObjectStream[LocalJet]
+        y: ObjectStream[LocalJet]
+
+    a = parse_as_ast(lambda e: nt_1(e.Jets(), e.Electrons()).x.Select(lambda j: j.pt())).body
+    a_resolved = resolve_syntatic_sugar(a)
+
+    a_expected = ast.parse("{'x': e.Jets(), 'y': e.Electrons()}.x.Select(lambda j: j.pt())")
+    assert ast.unparse(a_resolved) == ast.unparse(a_expected)
+
+
+def test_resolve_named_tuple_too_few_args():
+    "Make sure a named tuple becomes a dictionary"
+
+    nt_1 = namedtuple("nt_1", ["x", "y"])
+
+    a = parse_as_ast(lambda e: nt_1(e.Jets()).x.Select(lambda j: j.pt())).body  # type: ignore
+    a_resolved = resolve_syntatic_sugar(a)
+
+    a_expected = ast.parse("{'x': e.Jets()}.x.Select(lambda j: j.pt())")
+    assert ast.unparse(a_resolved) == ast.unparse(a_expected)
+
+
+def test_resolve_named_tuple_too_many_args():
+    "Make sure a named tuple becomes a dictionary"
+
+    nt_1 = namedtuple("nt_1", ["x", "y"])
+
+    with pytest.raises(ValueError):
+        a = parse_as_ast(
+            lambda e: nt_1(e.Jets(), e.Electrons(), e.Muons(), e.Jets()).x.Select(  # type: ignore
+                lambda j: j.pt()
+            )  # type: ignore
+        ).body
+        resolve_syntatic_sugar(a)
+
+
+def test_resolve_named_tuple_kwards():
+    "Make sure a named tuple becomes a dictionary"
+
+    nt_1 = namedtuple("nt_1", ["x", "y"])
+
+    a = parse_as_ast(lambda e: nt_1(x=e.Jets(), y=e.Electrons()).x.Select(lambda j: j.pt())).body
+    a_resolved = resolve_syntatic_sugar(a)
+
+    a_expected = ast.parse("{'x': e.Jets(), 'y': e.Electrons()}.x.Select(lambda j: j.pt())")
+    assert ast.unparse(a_resolved) == ast.unparse(a_expected)
+
+
+def test_resolve_random_class():
+    "A regular class should not be resolved"
+
+    class nt_1:
+        x: ObjectStream[LocalJet]
+        y: ObjectStream[LocalJet]
+
+    a = parse_as_ast(
+        lambda e: nt_1(x=e.Jets(), y=e.Electrons()).x.Select(lambda j: j.pt())  # type: ignore
+    ).body
+    a_resolved = resolve_syntatic_sugar(a)
+
+    assert "nt_1(" in ast.unparse(a_resolved)
diff --git a/tests/test_type_based_replacement.py b/tests/test_type_based_replacement.py
index fbdc0ab..861c7df 100644
--- a/tests/test_type_based_replacement.py
+++ b/tests/test_type_based_replacement.py
@@ -406,7 +406,7 @@ def test_collection_Custom_Method_default(caplog):
 
     @register_func_adl_os_collection
     class CustomCollection_default(ObjectStream[M]):
-        def __init__(self, a: ast.AST, item_type):
+        def __init__(self, a: ast.expr, item_type):
             super().__init__(a, item_type)
 
         def Take(self, n: int = 5) -> ObjectStream[M]: ...  # noqa
@@ -429,7 +429,7 @@ def test_collection_Custom_Method_Jet(caplog):
     M = TypeVar("M")
 
     class CustomCollection_Jet(ObjectStream[M]):
-        def __init__(self, a: ast.AST, item_type):
+        def __init__(self, a: ast.expr, item_type):
             super().__init__(a, item_type)
 
         def MyFirst(self) -> M: ...  # noqa
diff --git a/tests/test_util_ast.py b/tests/test_util_ast.py
index 476326f..0cf3085 100644
--- a/tests/test_util_ast.py
+++ b/tests/test_util_ast.py
@@ -1,5 +1,7 @@
 import ast
 import sys
+from collections import namedtuple
+from dataclasses import dataclass
 from enum import Enum
 from typing import Callable, cast
 
@@ -304,6 +306,34 @@ def test_parse_lambda_class_constant_in_module():
     assert ast.unparse(r) == ast.unparse(r_true)
 
 
+def test_parse_lambda_imported_class():
+    "Check that numpy and similar are properly passed"
+
+    import numpy as np
+
+    r = parse_as_ast(lambda e: np.cos(e))
+    assert "np.cos" in ast.unparse(r)
+
+
+def test_parse_dataclass_reference():
+    @dataclass
+    class my_data_class:
+        x: int
+
+    r = parse_as_ast(lambda e: my_data_class(x=e))
+
+    assert "<locals>.my_data_class" in ast.unparse(r)
+
+
+def test_parse_named_tuple_reference():
+
+    MyDataClass = namedtuple("MyDataClass", ["x"])
+
+    r = parse_as_ast(lambda e: MyDataClass(x=e))
+
+    assert "test_util_ast.MyDataClass" in ast.unparse(r)
+
+
 def test_parse_simple_func():
     "A oneline function defined at local scope"