Auto apply @dataclass to dataclass_array

PiperOrigin-RevId: 516194659

Author: Conchylicultor

CHANGELOG.md

@@ -26,6 +26,7 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 *   **Add `torch` support!**
 *   Add `.cpu()`, `.cuda()`, `.to()` methods to move the dataclass from
     devices when using torch.
+*   **Breaking**: `@dataclass(frozen=True)` is now automatically applied
 
 ## [1.3.0] - 2023-01-16
 

README.md

@@ -30,7 +30,6 @@ import dataclass_array as dca
 from dataclass_array.typing import FloatArray
 
 
-@dataclasses.dataclass(frozen=True)
 class Ray(dca.DataclassArray):
   pos: FloatArray['*batch_shape 3']
   dir: FloatArray['*batch_shape 3']
@@ -75,7 +74,6 @@ A `DataclassArray` has 2 types of fields:
     Static fields are also ignored in `jax.tree_map`.
 
 ```python
-@dataclasses.dataclass(frozen=True)
 class MyArray(dca.DataclassArray):
   # Array fields
   a: FloatArray['*batch_shape 3']  # Defined by `etils.array_types`
@@ -102,7 +100,6 @@ batching:
 2.  Decorate the method with `dca.vectorize_method`
 
 ```python
-@dataclasses.dataclass(frozen=True)
 class Camera(dca.DataclassArray):
   K: FloatArray['*batch_shape 4 4']
   resolution = tuple[int, int]
@@ -133,7 +130,7 @@ rays.shape == (num_cams, h, w)
 *   Instead of vectorizing a single axis, `@dca.vectorize_method` will vectorize
     over `*self.shape` (not just `self.shape[0]`). This is like if `vmap` was
     applied to `self.flatten()`
-*   When multiple arguments, axis with dimension `1` are brodcasted.
+*   When multiple arguments, axis with dimension `1` are broadcasted.
 
 For example, with `__matmul__(self, x: T) -> T`:
 

dataclass_array/array_dataclass.py

@@ -86,7 +86,6 @@ def dataclass_array(
 
   ```python
   @dca.dataclass_array()
-  @dataclasses.dataclass(frozen=True)
   class MyDataclass(dca.DataclassArray):
     ...
   ```
@@ -119,6 +118,31 @@ def decorator(cls):
   return decorator
 
 
+def array_field(
+    shape: Shape,
+    dtype: DTypeArg = float,
+    **field_kwargs,
+) -> dataclasses.Field[DcOrArray]:
+  """Dataclass array field.
+
+  See `dca.DataclassArray` for example.
+
+  Args:
+    shape: Inner shape of the field
+    dtype: Type of the field
+    **field_kwargs: Args forwarded to `dataclasses.field`
+
+  Returns:
+    The dataclass field.
+  """
+  # TODO(epot): Validate shape, dtype
+  dca_field = _ArrayFieldMetadata(
+      inner_shape_non_static=shape,
+      dtype=dtype,
+  )
+  return dataclasses.field(**field_kwargs, metadata={_METADATA_KEY: dca_field})
+
+
 class MetaDataclassArray(type):
   """DataclassArray metaclass."""
 
@@ -128,13 +152,21 @@ def __getitem__(cls, spec):
     return Annotated[cls, field_utils.ShapeAnnotation(spec)]
 
 
+@typing_extensions.dataclass_transform(  # pytype: disable=not-supported-yet
+    kw_only_default=True,
+    # TODO(b/272524683):Restore field specifier
+    # field_specifiers=(
+    #     dataclasses.Field,
+    #     dataclasses.field,
+    #     array_field,
+    # ),
+)
 class DataclassArray(metaclass=MetaDataclassArray):
   """Dataclass which behaves like an array.
 
   Usage:
 
   ```python
-  @dataclasses.dataclass
   class Square(DataclassArray):
     pos: f32['*shape 2']
     scale: f32['*shape']
@@ -179,7 +211,6 @@ class Square(DataclassArray):
 
   Field which do not satisfy any of the above conditions are static (including
   field annotated with `field: np.ndarray` or similar).
-
   """
 
   # Child class inherit the default params by default, but can also
@@ -194,8 +225,21 @@ class Square(DataclassArray):
   _shape: Shape
   _xnp: enp.NpModule
 
-  def __init_subclass__(cls, **kwargs):
+  def __init_subclass__(
+      cls,
+      frozen=True,
+      **kwargs,
+  ):
     super().__init_subclass__(**kwargs)
+
+    if not frozen:
+      raise ValueError(f'{cls} cannot be `frozen=False`.')
+
+    # Apply dataclass (in-place)
+    if not typing.TYPE_CHECKING:
+      # TODO(b/227290126): Create pytype issues
+      dataclasses.dataclass(frozen=True)(cls)
+
     # TODO(epot): Could have smart __repr__ which display types if array have
     # too many values (maybe directly in `edc.field(repr=...)`).
     edc.dataclass(kw_only=True, repr=True, auto_cast=False)(cls)
@@ -212,6 +256,11 @@ def __init_subclass__(cls, **kwargs):
     # `__dca_non_init_fields__` (fields should be merged from `.mro()`)
     cls.__dca_non_init_fields__ = set(cls.__dca_non_init_fields__)
 
+  if typing.TYPE_CHECKING:
+    # TODO(b/242839979): pytype do not support PEP 681 -- Data Class Transforms
+    def __init__(self, **kwargs):
+      pass
+
   def __post_init__(self) -> None:
     """Validate and normalize inputs."""
     cls = type(self)
@@ -755,7 +804,6 @@ def _init_cls(self: DataclassArray) -> None:
 
   This will:
 
-  * Validate the `@dataclass(frozen=True)` is correctly applied
   * Extract the types annotations, detect which fields are arrays or static,
     and store the result in `_dca_fields_metadata`
   * For static `DataclassArray` (class with only static fields), it will
@@ -767,12 +815,6 @@ def _init_cls(self: DataclassArray) -> None:
   """
   cls = type(self)
 
-  # Make sure the dataclass was registered and frozen
-  if not dataclasses.is_dataclass(cls) or not cls.__dataclass_params__.frozen:  # pytype: disable=attribute-error
-    raise ValueError(
-        '`dca.DataclassArray` need to be @dataclasses.dataclass(frozen=True)'
-    )
-
   # The first time, compute typing annotations & metadata
   # At this point, `ForwardRef` should have been resolved.
   try:
@@ -797,7 +839,7 @@ def _init_cls(self: DataclassArray) -> None:
     )
 
   dca_fields_metadata = {
-      f.name: _make_field_metadata(f, hints) for f in dataclasses.fields(cls)
+      f.name: _make_field_metadata(f, hints) for f in dataclasses.fields(cls)  # pytype: disable=wrong-arg-types
   }
   dca_fields_metadata = {  # Filter `None` values (static fields)
       k: v for k, v in dca_fields_metadata.items() if v is not None
@@ -909,31 +951,6 @@ class _TreeMetadata:
   non_array_field_kwargs: dict[str, Any]
 
 
-def array_field(
-    shape: Shape,
-    dtype: DTypeArg = float,
-    **field_kwargs,
-) -> dataclasses.Field[DcOrArray]:
-  """Dataclass array field.
-
-  See `dca.DataclassArray` for example.
-
-  Args:
-    shape: Inner shape of the field
-    dtype: Type of the field
-    **field_kwargs: Args forwarded to `dataclasses.field`
-
-  Returns:
-    The dataclass field.
-  """
-  # TODO(epot): Validate shape, dtype
-  dca_field = _ArrayFieldMetadata(
-      inner_shape_non_static=shape,
-      dtype=dtype,
-  )
-  return dataclasses.field(**field_kwargs, metadata={_METADATA_KEY: dca_field})
-
-
 # TODO(epot): Should refactor `_ArrayField` in `_DataclassArrayField` and
 # `_ArrayField` depending on whether dtype is `DataclassArray` or not.
 # Alternativelly, maybe should create a `DcArrayDType` dtype instead.

dataclass_array/array_dataclass_test.py

@@ -16,7 +16,6 @@
 
 from __future__ import annotations
 
-import dataclasses
 from typing import Any
 
 import dataclass_array as dca
@@ -49,7 +48,6 @@ def assert_val(
 
 
 @dca.dataclass_array(broadcast=True, cast_dtype=True)
-@dataclasses.dataclass(frozen=True)
 class Point(dca.DataclassArray):
   x: f32['*shape']
   y: f32['*shape']
@@ -77,7 +75,6 @@ def assert_val(p: Point, shape: Shape, xnp: enp.NpModule = None):
 
 
 @dca.dataclass_array(broadcast=True, cast_dtype=True)
-@dataclasses.dataclass(frozen=True)
 class Isometrie(dca.DataclassArray):
   r: f32['... 3 3']
   t: i32[..., 2]
@@ -105,7 +102,6 @@ def assert_val(p: Isometrie, shape: Shape, xnp: enp.NpModule = None):
 
 
 @dca.dataclass_array(broadcast=True, cast_dtype=True)
-@dataclasses.dataclass(frozen=True)
 class Nested(dca.DataclassArray):
   # pytype: disable=annotation-type-mismatch
   iso: Isometrie
@@ -142,7 +138,6 @@ def assert_val(p: Nested, shape: Shape, xnp: enp.NpModule = None):
     Isometrie.assert_val(p.iso_batched, shape=shape + (3, 7), xnp=xnp)
 
 
-@dataclasses.dataclass(frozen=True)
 class OnlyStatic(dca.DataclassArray):
   """Dataclass with no array fields."""
 
@@ -172,7 +167,6 @@ def assert_val(p: OnlyStatic, shape: Shape, xnp: enp.NpModule = None):
     assert p.y == 1
 
 
-@dataclasses.dataclass(frozen=True)
 class NestedOnlyStatic(dca.DataclassArray):
   """Dataclass with only nested array fields."""
 
@@ -198,7 +192,6 @@ def assert_val(p: NestedOnlyStatic, shape: Shape, xnp: enp.NpModule = None):
 
 
 @dca.dataclass_array(broadcast=True, cast_dtype=True)
-@dataclasses.dataclass(frozen=True)
 class WithStatic(dca.DataclassArray):
   """Mix of static and array fields."""
 
@@ -413,7 +406,6 @@ def test_wrong_input_type():
     )
 
   @dca.dataclass_array(broadcast=True, cast_dtype=True)
-  @dataclasses.dataclass(frozen=True)
   class PointWrapper(dca.DataclassArray):
     pts: Point
     rgb: f32['*shape 3']
@@ -642,7 +634,6 @@ def fn(p: WithStatic) -> WithStatic:
 
 @enp.testing.parametrize_xnp()
 def test_dataclass_params_no_cast(xnp: enp.NpModule):
-  @dataclasses.dataclass(frozen=True)
   class PointNoCast(dca.DataclassArray):
     x: FloatArray['*shape']
     y: IntArray['*shape']
@@ -665,7 +656,6 @@ class PointNoCast(dca.DataclassArray):
 @enp.testing.parametrize_xnp()
 def test_dataclass_params_no_list(xnp: enp.NpModule):
   @dca.dataclass_array(cast_list=False)
-  @dataclasses.dataclass(frozen=True)
   class PointNoList(dca.DataclassArray):
     x: FloatArray['*shape']
     y: IntArray['*shape']
@@ -679,7 +669,6 @@ class PointNoList(dca.DataclassArray):
 
 @enp.testing.parametrize_xnp()
 def test_dataclass_params_no_broadcast(xnp: enp.NpModule):
-  @dataclasses.dataclass(frozen=True)
   class PointNoBroadcast(dca.DataclassArray):
     x: FloatArray['*shape']
     y: IntArray['*shape']
@@ -694,7 +683,6 @@ class PointNoBroadcast(dca.DataclassArray):
 @enp.testing.parametrize_xnp()
 @pytest.mark.parametrize('batch_shape', [(), (1, 3)])
 def test_dataclass_none_shape(xnp: enp.NpModule, batch_shape: Shape):
-  @dataclasses.dataclass(frozen=True)
   class PointDynamicShape(dca.DataclassArray):
     x: FloatArray[..., None, None]
     y: IntArray['... 3 _']

dataclass_array/import_test.py

@@ -20,14 +20,12 @@
 
 from __future__ import annotations
 
-import dataclasses
 import sys
 
 import dataclass_array as dca
 from etils import enp
 
 
-@dataclasses.dataclass(frozen=True)
 class A(dca.DataclassArray):
   x: dca.typing.f32['*s']
 

dataclass_array/testing.py

@@ -16,7 +16,6 @@
 
 from __future__ import annotations
 
-import dataclasses
 import functools
 from typing import Any, Optional
 
@@ -27,7 +26,6 @@
 import numpy as np
 
 
-@dataclasses.dataclass(frozen=True)
 class Ray(array_dataclass.DataclassArray):
   """Dummy dataclass array for testing."""
 

dataclass_array/type_parsing_test.py

@@ -16,7 +16,6 @@
 
 from __future__ import annotations
 
-import dataclasses
 from typing import Optional, List, Union
 
 import dataclass_array as dca
@@ -28,7 +27,6 @@
 Ray = dca.testing.Ray
 
 
-@dataclasses.dataclass(frozen=True)
 class Camera(dca.DataclassArray):
   pos: FloatArray[..., 3]
   dir: FloatArray[..., 3]

dataclass_array/vectorization.py

@@ -100,7 +100,6 @@ def vectorize_method(
   Example:
 
   ```
-  @dataclasses.dataclass(frozen=True)
   class Point3d(dca.DataclassArray):
     p: f32['*shape 3']
 

dataclass_array/vectorization_test.py

@@ -150,7 +150,6 @@ def fn(self, arg):
 def test_replace_dca(xnp: enp.NpModule):
   # Ensure that the non-init static fields are correctly forwarded.
 
-  @dataclasses.dataclass(frozen=True)
   class DataclassWithNonInit(dca.DataclassArray):
     """Dataclass with a non-init (static) field."""