Add named dataclasses as a feature

func_adl/util_ast.py

@@ -5,6 +5,7 @@
 import sys
 import tokenize
 from collections import defaultdict
+from dataclasses import is_dataclass
 from types import ModuleType
 from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union, cast
 
@@ -328,6 +329,10 @@ def visit_Name(self, node: ast.Name) -> Any:
                 # If it is something we know how to make into a literal, we just send it down
                 # like that.
                 return as_literal(v)
+
+            if is_dataclass(v):
+                return ast.Constant(value=v, kind=None)
+
         return node
 
     def visit_Lambda(self, node: ast.Lambda) -> Any:

tests/test_type_based_replacement.py

@@ -1,6 +1,7 @@
 import ast
 import copy
 import logging
+from dataclasses import dataclass
 from typing import Any, Callable, Iterable, Optional, Tuple, Type, TypeVar, cast
 
 import pytest
@@ -14,6 +15,7 @@
     remap_by_types,
     remap_from_lambda,
 )
+from func_adl.util_ast import parse_as_ast
 
 
 class Track:
@@ -561,6 +563,57 @@ def test_indexed_tuple_bad_slice():
     assert "is not valid" in str(e)
 
 
+def test_dataclass_create():
+    "Check a simple, correct, data class reference"
+
+    @dataclass
+    class MyData:
+        pt: float
+        eta: float
+
+    s = parse_as_ast(lambda e: e.Jets().Select(lambda j: MyData(j.pt(), j.eta()))).body
+    objs = ObjectStream[Event](ast.Name(id="e", ctx=ast.Load()))
+
+    new_objs, new_s, expr_type = remap_by_types(objs, "e", Event, s)
+
+    assert expr_type == Iterable[MyData]
+
+
+def test_dataclass_simple_reference():
+    "Check a simple, correct, data class reference"
+
+    @dataclass
+    class MyData:
+        pt: float
+        eta: float
+
+    s = ast_lambda("e.Jets().Select(lambda j: MyData(j.pt(), j.eta())).Select(lambda d: d.pt)")
+
+    objs = ObjectStream[Event](ast.Name(id="e", ctx=ast.Load()))
+
+    new_objs, new_s, expr_type = remap_by_types(objs, "e", Event, s)
+
+    assert expr_type == Iterable[MyData]
+
+
+def test_dataclass_simple_bad_reference():
+    "Check a simple, correct, data class reference"
+
+    @dataclass
+    class MyData:
+        pt: float
+        eta: float
+
+    s = ast_lambda("e.Jets().Select(lambda j: MyData(j.pt(), j.eta())).Select(lambda d: d.bad_pt)")
+
+    objs = ObjectStream[Event](ast.Name(id="e", ctx=ast.Load()))
+
+    with pytest.raises(ValueError) as e:
+        remap_by_types(objs, "e", Event, s)
+
+    assert "bad_pt" in str(e)
+
+
 def test_collection_Select_meta(caplog):
     "A simple collection"
     caplog.set_level(logging.WARNING)

tests/test_util_ast.py

@@ -1,5 +1,6 @@
 import ast
 import sys
+from dataclasses import dataclass
 from typing import Callable, cast
 
 import pytest
@@ -858,6 +859,17 @@ def __init__(self, n):
     )
 
 
+def test_dataclass_creation():
+    @dataclass
+    class my_data:
+        pt: float
+        eta: int
+
+    a = parse_as_ast(lambda e: my_data(10.0, 20))
+    d_text = ast.dump(a)
+    assert "Constant(value=<class 'tests" in d_text
+
+
 def test_parse_metadata_there():
     recoreded = None