diff --git a/.gitignore b/.gitignore index d53bcd08..b8cb797f 100644 --- a/.gitignore +++ b/.gitignore @@ -121,3 +121,6 @@ lightning_logs Formatting .ruff_cache/ + +# do not publish to PyPI +uv.lock diff --git a/CHANGELOG.md b/CHANGELOG.md index 0e5a51d5..ce29fc20 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - Added support for Python 3.14 ([#464](https://github.com/Lightning-AI/utilities/pull/464)) +- Added `apply_to_collection`(s) support for frozen dataclasses ([#499](https://github.com/Lightning-AI/utilities/pull/499)) ### Changed diff --git a/src/lightning_utilities/core/apply_func.py b/src/lightning_utilities/core/apply_func.py index 4303fc2a..a851bb62 100644 --- a/src/lightning_utilities/core/apply_func.py +++ b/src/lightning_utilities/core/apply_func.py @@ -26,6 +26,55 @@ def is_dataclass_instance(obj: object) -> bool: return dataclasses.is_dataclass(obj) and not isinstance(obj, type) +def is_dataclass_frozen(obj: object) -> bool: + """Return True if the given object is a frozen dataclass instance.""" + # https://stackoverflow.com/a/75958306 + return getattr(getattr(obj, "__dataclass_params__", None), "frozen", False) + + +def _dataclass_has_init_vars(data: Any) -> bool: + """Return True if the dataclass declares any ``InitVar`` (pseudo-)fields. + + ``InitVar`` fields are excluded from ``dataclasses.fields`` and their values are consumed by ``__post_init__`` + rather than stored on the instance, so a frozen dataclass that declares them cannot be reconstructed from its stored + fields. + + """ + return any(f._field_type is dataclasses._FIELD_INITVAR for f in data.__dataclass_fields__.values()) # type: ignore[attr-defined] + + +def _reconstruct_frozen_dataclass(data: Any, apply_field: Callable[[Any], Any]) -> Any: + """Reconstruct a frozen dataclass by applying ``apply_field`` to each ``init`` field's value. + + Frozen dataclasses cannot be mutated in place, so instead of the deepcopy-and-``setattr`` approach + used for mutable dataclasses, a new instance is built from the (transformed) init fields. Fields + with ``init=False`` are retained from the original instance and written back via + ``object.__setattr__`` (this overrides whatever ``__post_init__`` derived during construction, + keeping parity with the retain-old behavior of the mutable path). + + Raises: + ValueError: If the frozen dataclass declares ``InitVar`` fields, which cannot be reconstructed. + + """ + if _dataclass_has_init_vars(data): + raise ValueError( + "A frozen dataclass with `InitVar` fields was passed to `apply_to_collection`(s)" + " but this is not supported: such instances cannot be reconstructed." + ) + init_fields = {} + non_init_fields = {} + for field in dataclasses.fields(data): + field_value = getattr(data, field.name) + if field.init: + init_fields[field.name] = apply_field(field_value) + else: + non_init_fields[field.name] = field_value + result = type(data)(**init_fields) + for field_name, field_value in non_init_fields.items(): + object.__setattr__(result, field_name, field_value) + return result + + def apply_to_collection( data: Any, dtype: type | Any | tuple[type | Any], @@ -33,7 +82,6 @@ def apply_to_collection( *args: Any, wrong_dtype: type | tuple[type, ...] | None = None, include_none: bool = True, - allow_frozen: bool = False, **kwargs: Any, ) -> Any: """Recursively applies a function to all elements of a certain dtype. @@ -46,14 +94,13 @@ def apply_to_collection( wrong_dtype: the given function won't be applied if this type is specified and the given collections is of the ``wrong_dtype`` even if it is of type ``dtype`` include_none: Whether to include an element if the output of ``function`` is ``None``. - allow_frozen: Whether not to error upon encountering a frozen dataclass instance. **kwargs: keyword arguments (will be forwarded to calls of ``function``) Returns: The resulting collection """ - if include_none is False or wrong_dtype is not None or allow_frozen is True: + if include_none is False or wrong_dtype is not None: # not worth implementing these on the fast path: go with the slower option return _apply_to_collection_slow( data, @@ -62,7 +109,6 @@ def apply_to_collection( *args, wrong_dtype=wrong_dtype, include_none=include_none, - allow_frozen=allow_frozen, **kwargs, ) # fast path for the most common cases: @@ -82,7 +128,6 @@ def apply_to_collection( *args, wrong_dtype=wrong_dtype, include_none=include_none, - allow_frozen=allow_frozen, **kwargs, ) @@ -94,7 +139,6 @@ def _apply_to_collection_slow( *args: Any, wrong_dtype: type | tuple[type, ...] | None = None, include_none: bool = True, - allow_frozen: bool = False, **kwargs: Any, ) -> Any: # Breaking condition @@ -105,6 +149,25 @@ def _apply_to_collection_slow( # Recursively apply to collection items if is_dataclass_instance(data): + if is_dataclass_frozen(data): + # frozen dataclass: reconstruct a new instance from (transformed) init fields + def _apply_field(field_value: Any) -> Any: + v = _apply_to_collection_slow( + field_value, + dtype, + function, + *args, + wrong_dtype=wrong_dtype, + include_none=include_none, + **kwargs, + ) + if not include_none and v is None: # retain old value + v = field_value + return v + + return _reconstruct_frozen_dataclass(data, _apply_field) + + # mutable dataclass: deepcopy + setattr (handles InitVar, init=False, and cached_property reset) # make a deepcopy of the data, # but do not deepcopy mapped fields since the computation would # be wasted on values that likely get immediately overwritten @@ -126,24 +189,16 @@ def _apply_to_collection_slow( *args, wrong_dtype=wrong_dtype, include_none=include_none, - allow_frozen=allow_frozen, **kwargs, ) if not field_init or (not include_none and v is None): # retain old value v = getattr(data, field_name) - try: - setattr(result, field_name, v) - except dataclasses.FrozenInstanceError as e: - if allow_frozen: - # Quit early if we encounter a frozen data class; return `result` as is. - break - raise ValueError( - "A frozen dataclass was passed to `apply_to_collection` but this is not allowed." - ) from e + setattr(result, field_name, v) # Explicitly resetting cached property. for cached_name in filter( - lambda k: isinstance(getattr(type(data), k), cached_property), vars(type(data)).keys() + lambda k: isinstance(getattr(type(data), k), cached_property), + vars(type(data)).keys(), ): vars(result).pop(cached_name, None) return result @@ -158,7 +213,6 @@ def _apply_to_collection_slow( *args, wrong_dtype=wrong_dtype, include_none=include_none, - allow_frozen=allow_frozen, **kwargs, ) if include_none or v is not None: @@ -179,7 +233,6 @@ def _apply_to_collection_slow( *args, wrong_dtype=wrong_dtype, include_none=include_none, - allow_frozen=allow_frozen, **kwargs, ) if include_none or v is not None: @@ -214,7 +267,8 @@ def apply_to_collections( A collection with the same structure as the input where matching elements are transformed. Raises: - ValueError: If sequence collections have different sizes. + ValueError: If sequence collections have different sizes, or if a frozen dataclass with ``InitVar`` + fields is passed. TypeError: If dataclass inputs are mismatched (different types or fields), or if ``data1`` is a dataclass instance but ``data2`` is not. @@ -257,9 +311,36 @@ def apply_to_collections( ) if not ( len(dataclasses.fields(data1)) == len(dataclasses.fields(data2)) - and all(map(lambda f1, f2: isinstance(f1, type(f2)), dataclasses.fields(data1), dataclasses.fields(data2))) + and all( + map( + lambda f1, f2: isinstance(f1, type(f2)), + dataclasses.fields(data1), + dataclasses.fields(data2), + ) + ) ): raise TypeError("Dataclasses fields do not match.") + + if is_dataclass_frozen(data1): + # frozen dataclass: reconstruct from the zipped (transformed) init fields. + # build a lookup of data2's fields so we can zip by position during reconstruction. + fields2 = {field.name: getattr(data2, field.name) for field in dataclasses.fields(data2)} + field_iter = iter(dataclasses.fields(data1)) + + def _apply_field(field_value1: Any) -> Any: + field = next(field_iter) + return apply_to_collections( + field_value1, + fields2[field.name], + dtype, + function, + *args, + wrong_dtype=wrong_dtype, + **kwargs, + ) + + return _reconstruct_frozen_dataclass(data1, _apply_field) + # make a deepcopy of the data, # but do not deepcopy mapped fields since the computation would # be wasted on values that likely get immediately overwritten @@ -276,9 +357,10 @@ def apply_to_collections( result = deepcopy(data1, memo=memo) # apply function to each field - for (field_name, (field_value1, field_init1)), (_, (field_value2, field_init2)) in zip( - fields[0].items(), fields[1].items() - ): + for (field_name, (field_value1, field_init1)), ( + _, + (field_value2, field_init2), + ) in zip(fields[0].items(), fields[1].items()): v = None if field_init1 and field_init2: v = apply_to_collections( @@ -292,12 +374,7 @@ def apply_to_collections( ) if not field_init1 or not field_init2 or v is None: # retain old value return apply_to_collection(data1, dtype, function, *args, wrong_dtype=wrong_dtype, **kwargs) - try: - setattr(result, field_name, v) - except dataclasses.FrozenInstanceError as e: - raise ValueError( - "A frozen dataclass was passed to `apply_to_collections` but this is not allowed." - ) from e + setattr(result, field_name, v) return result return apply_to_collection(data1, dtype, function, *args, wrong_dtype=wrong_dtype, **kwargs) diff --git a/tests/unittests/core/test_apply_func.py b/tests/unittests/core/test_apply_func.py index 3860b942..e7578f99 100644 --- a/tests/unittests/core/test_apply_func.py +++ b/tests/unittests/core/test_apply_func.py @@ -349,18 +349,132 @@ class Foo: var: int foo = Foo(0) - with pytest.raises(ValueError, match="frozen dataclass was passed"): + # frozen dataclasses are now reconstructed (a new transformed instance is returned) + result = apply_to_collection(foo, int, lambda x: x + 1) + assert result == Foo(1) + # the original instance is left untouched + assert foo == Foo(0) + + +def test_apply_to_collection_frozen_dataclass_nested(): + @dataclasses.dataclass(frozen=True) + class Inner: + val: torch.Tensor + + @dataclasses.dataclass(frozen=True) + class Outer: + inner: Inner + label: torch.Tensor + + outer = Outer(inner=Inner(val=torch.tensor(1.0)), label=torch.tensor(2.0)) + result = apply_to_collection(outer, torch.Tensor, lambda x: x * 3) + assert torch.equal(result.inner.val, torch.tensor(3.0)) + assert torch.equal(result.label, torch.tensor(6.0)) + # original untouched + assert torch.equal(outer.inner.val, torch.tensor(1.0)) + + +def test_apply_to_collection_frozen_dataclass_with_non_init_field(): + @dataclasses.dataclass(frozen=True) + class Foo: + var: int + computed: int = dataclasses.field(init=False) + + def __post_init__(self): + object.__setattr__(self, "computed", self.var * 2) + + foo = Foo(5) + result = apply_to_collection(foo, int, lambda x: x + 1) + assert result.var == 6 + # init=False fields are retained from the original (parity with the mutable path), + # so `computed` keeps its original value (10) rather than being recomputed from the new var (12) + assert result.computed == 10 + # original untouched + assert foo.var == 5 + assert foo.computed == 10 + + +def test_apply_to_collection_frozen_dataclass_with_class_var(): + @dataclasses.dataclass(frozen=True) + class Foo: + const: ClassVar[int] = 7 + var: int + + foo = Foo(5) + # a ClassVar must not be mistaken for an InitVar; reconstruction should succeed + result = apply_to_collection(foo, int, lambda x: x + 1) + assert result.var == 6 + assert result.const == 7 + assert Foo.const == 7 + assert foo.var == 5 + + +def test_apply_to_collection_frozen_dataclass_with_init_var_raises(): + @dataclasses.dataclass(frozen=True) + class Foo: + var: int + scale: InitVar[int] = 1 + + def __post_init__(self, scale): + object.__setattr__(self, "var", self.var * scale) + + foo = Foo(5, scale=2) + # frozen dataclasses with InitVar fields cannot be reconstructed + with pytest.raises(ValueError, match="InitVar"): apply_to_collection(foo, int, lambda x: x + 1) -def test_apply_to_collection_allow_frozen_dataclass(): +def test_apply_to_collections_frozen_dataclass(): + @dataclasses.dataclass(frozen=True) + class Foo: + a: torch.Tensor + b: torch.Tensor + + def __eq__(self, o: object) -> bool: + """Equal.""" + if not isinstance(o, Foo): + return NotImplemented + return torch.equal(self.a, o.a) and torch.equal(self.b, o.b) + + f1 = Foo(torch.tensor([1.0]), torch.tensor([2.0])) + f2 = Foo(torch.tensor([10.0]), torch.tensor([20.0])) + result = apply_to_collections(f1, f2, torch.Tensor, lambda x, y: x + y) + assert torch.equal(result.a, torch.tensor([11.0])) + assert torch.equal(result.b, torch.tensor([22.0])) + # original untouched + assert torch.equal(f1.a, torch.tensor([1.0])) + + +def test_apply_to_collections_frozen_dataclass_nested(): + @dataclasses.dataclass(frozen=True) + class Inner: + v: torch.Tensor + + @dataclasses.dataclass(frozen=True) + class Outer: + inner: Inner + x: torch.Tensor + + o1 = Outer(Inner(torch.tensor([1.0])), torch.tensor([2.0])) + o2 = Outer(Inner(torch.tensor([10.0])), torch.tensor([20.0])) + result = apply_to_collections(o1, o2, torch.Tensor, lambda a, b: a + b) + assert torch.equal(result.inner.v, torch.tensor([11.0])) + assert torch.equal(result.x, torch.tensor([22.0])) + + +def test_apply_to_collections_frozen_dataclass_with_init_var_raises(): @dataclasses.dataclass(frozen=True) class Foo: var: int + scale: InitVar[int] = 1 - foo = Foo(0) - result = apply_to_collection(foo, int, lambda x: x + 1, allow_frozen=True) - assert foo == result + def __post_init__(self, scale): + object.__setattr__(self, "var", self.var * scale) + + f1 = Foo(5, scale=2) + f2 = Foo(5, scale=2) + with pytest.raises(ValueError, match="InitVar"): + apply_to_collections(f1, f2, int, lambda a, b: a + b) def test_apply_to_collection_with_cached_property_dataclass():