allow to exclude fields in dataclass_array_container

Author: alexfikl

arraycontext/__init__.py

@@ -44,7 +44,7 @@
         serialize_container, deserialize_container,
         register_multivector_as_array_container)
 from .container.arithmetic import with_container_arithmetic
-from .container.dataclass import dataclass_array_container
+from .container.dataclass import dataclass_array_container, ExcludedField
 
 from .container.traversal import (
         map_array_container,
@@ -85,7 +85,7 @@
         "serialize_container", "deserialize_container",
         "register_multivector_as_array_container",
         "with_container_arithmetic",
-        "dataclass_array_container",
+        "dataclass_array_container", "ExcludedField",
 
         "map_array_container", "multimap_array_container",
         "rec_map_array_container", "rec_multimap_array_container",

arraycontext/container/dataclass.py

@@ -2,6 +2,8 @@
 
 """
 .. currentmodule:: arraycontext
+
+.. autoclass:: ExcludedField
 .. autofunction:: dataclass_array_container
 """
 
@@ -31,21 +33,49 @@
 """
 
 from dataclasses import fields
+
+try:
+    # NOTE: mypy fails with `Module "typing" has no attribute "get_args"`
+    from typing import (                                # type: ignore[attr-defined]
+            _AnnotatedAlias, _GenericAlias, get_args)
+except ImportError:
+    from typing_extensions import (                     # type: ignore[attr-defined]
+            _AnnotatedAlias, _GenericAlias, get_args)
+
 from arraycontext.container import is_array_container_type
 
 
 # {{{ dataclass containers
 
+class ExcludedField:
+    """Can be used to annotate dataclass fields to be excluded from the container.
+
+    This can be done using :class:`typing.Annotated` as follows
+
+    .. code:: python
+
+        @dataclass
+        class MyClass:
+            x: np.ndarray
+            y: Annotated[np.ndarray, ExcludedField]
+    """
+
+
 def dataclass_array_container(cls: type) -> type:
-    """A class decorator that makes the class to which it is applied a
+    """A class decorator that makes the class to which it is applied an
     :class:`ArrayContainer` by registering appropriate implementations of
     :func:`serialize_container` and :func:`deserialize_container`.
     *cls* must be a :func:`~dataclasses.dataclass`.
 
     Attributes that are not array containers are allowed. In order to decide
     whether an attribute is an array container, the declared attribute type
     is checked by the criteria from :func:`is_array_container_type`.
+
+    To explicitly exclude fields from the container serialization (that would
+    otherwise be recognized as array containers), use :class:`typing.Annotated`
+    using :class:`ExcludedField`.
     """
+
     from dataclasses import is_dataclass, Field
     assert is_dataclass(cls)
 
@@ -59,8 +89,12 @@ def is_array_field(f: Field) -> bool:
                 raise TypeError(
                         f"string annotation on field '{f.name}' not supported")
 
-            from typing import _SpecialForm
-            if isinstance(f.type, _SpecialForm):
+            # FIXME: is there a nicer way to recognize that we hit Annotated?
+            if isinstance(f.type, _AnnotatedAlias):
+                if any(arg is ExcludedField for arg in get_args(f.type)):
+                    return False
+
+            if isinstance(f.type, _GenericAlias):
                 raise TypeError(
                         f"typing annotation not supported on field '{f.name}': "
                         f"'{f.type!r}'")

test/test_utils.py

@@ -93,6 +93,33 @@ class ArrayContainerWithInitFalse:
 # }}}
 
 
+# {{{ test_dataclass_excluded_fields
+
+def test_dataclass_excluded_fields():
+    from dataclasses import dataclass
+    from arraycontext import dataclass_array_container, ExcludedField
+
+    try:
+        from typing import Annotated
+    except ImportError:
+        from typing_extensions import Annotated
+
+    @dataclass_array_container
+    @dataclass(frozen=True)
+    class ExcludedFoo:
+        x: np.ndarray
+        y: np.ndarray
+        excluded: Annotated[np.ndarray, ExcludedField]
+
+    ary = np.array([42], dtype=object)
+    x = ExcludedFoo(x=ary, y=ary, excluded=ary)
+
+    from arraycontext import serialize_container
+    assert len(serialize_container(x)) == 2
+
+# }}}
+
+
 if __name__ == "__main__":
     import sys
     if len(sys.argv) > 1: