From 1f18dba2ac0935925549c6b27775eab2041974c5 Mon Sep 17 00:00:00 2001 From: deependujha Date: Sun, 28 Jun 2026 15:12:49 +0530 Subject: [PATCH 1/7] feat: support reconstruction of frozen dataclass instances in apply_to_collection --- .gitignore | 3 + src/lightning_utilities/core/apply_func.py | 131 ++++++++++++++++---- tests/unittests/core/test_apply_func.py | 137 ++++++++++++++++++++- 3 files changed, 241 insertions(+), 30 deletions(-) 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/src/lightning_utilities/core/apply_func.py b/src/lightning_utilities/core/apply_func.py index 4303fc2a..a9742b08 100644 --- a/src/lightning_utilities/core/apply_func.py +++ b/src/lightning_utilities/core/apply_func.py @@ -3,6 +3,7 @@ # http://www.apache.org/licenses/LICENSE-2.0 # import dataclasses +import warnings from collections import OrderedDict, defaultdict from collections.abc import Callable, Mapping, Sequence from copy import deepcopy @@ -26,6 +27,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()) + + +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 +83,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 +95,24 @@ 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 "allow_frozen" in kwargs: + # `allow_frozen` was removed: frozen dataclasses are now reconstructed automatically. + # Drop it here so it is not forwarded to `function`, and warn that it is a no-op. + kwargs.pop("allow_frozen") + warnings.warn( + "The `allow_frozen` argument of `apply_to_collection` is deprecated and has no effect:" + " frozen dataclasses are now reconstructed automatically. It will be removed in a future release.", + DeprecationWarning, + stacklevel=2, + ) + + 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 +121,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 +140,6 @@ def apply_to_collection( *args, wrong_dtype=wrong_dtype, include_none=include_none, - allow_frozen=allow_frozen, **kwargs, ) @@ -94,7 +151,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 +161,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,20 +201,11 @@ 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( @@ -158,7 +224,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 +244,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 +278,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. @@ -260,6 +325,27 @@ def apply_to_collections( 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 @@ -292,12 +378,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..8264ea2e 100644 --- a/tests/unittests/core/test_apply_func.py +++ b/tests/unittests/core/test_apply_func.py @@ -349,18 +349,145 @@ class Foo: var: int foo = Foo(0) - with pytest.raises(ValueError, match="frozen dataclass was passed"): - apply_to_collection(foo, int, lambda x: x + 1) + # 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_allow_frozen_dataclass(): +def test_apply_to_collection_allow_frozen_kwarg_is_ignored(): @dataclasses.dataclass(frozen=True) class Foo: var: int foo = Foo(0) - result = apply_to_collection(foo, int, lambda x: x + 1, allow_frozen=True) - assert foo == result + # `allow_frozen` is deprecated; passing it must warn, must not error, and must not be forwarded to `function` + with pytest.warns(DeprecationWarning, match="allow_frozen"): + result = apply_to_collection(foo, int, lambda x: x + 1, allow_frozen=True) + assert result == Foo(1) + 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_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 + + 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(): From 528e524c881f61432f38e4b9fb3796f9ee9b954e Mon Sep 17 00:00:00 2001 From: deependujha Date: Sun, 28 Jun 2026 15:21:43 +0530 Subject: [PATCH 2/7] update --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0e5a51d5..e81f1aef 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_collections` support for frozen dataclasses with `InitVar` ([#499](https://github.com/Lightning-AI/utilities/pull/499)) ### Changed From 7936865d91c9e7314c4f1d44922588f2c46fa645 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 28 Jun 2026 09:53:18 +0000 Subject: [PATCH 3/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/lightning_utilities/core/apply_func.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/lightning_utilities/core/apply_func.py b/src/lightning_utilities/core/apply_func.py index a9742b08..04538f85 100644 --- a/src/lightning_utilities/core/apply_func.py +++ b/src/lightning_utilities/core/apply_func.py @@ -36,9 +36,9 @@ def is_dataclass_frozen(obj: object) -> bool: 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. + ``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()) From d0d2169f101ee9f2b5950074942f332356cc0d44 Mon Sep 17 00:00:00 2001 From: deependujha Date: Sun, 28 Jun 2026 15:25:27 +0530 Subject: [PATCH 4/7] update --- src/lightning_utilities/core/apply_func.py | 11 ----------- tests/unittests/core/test_apply_func.py | 13 ------------- 2 files changed, 24 deletions(-) diff --git a/src/lightning_utilities/core/apply_func.py b/src/lightning_utilities/core/apply_func.py index 04538f85..1eea29a6 100644 --- a/src/lightning_utilities/core/apply_func.py +++ b/src/lightning_utilities/core/apply_func.py @@ -101,17 +101,6 @@ def apply_to_collection( The resulting collection """ - if "allow_frozen" in kwargs: - # `allow_frozen` was removed: frozen dataclasses are now reconstructed automatically. - # Drop it here so it is not forwarded to `function`, and warn that it is a no-op. - kwargs.pop("allow_frozen") - warnings.warn( - "The `allow_frozen` argument of `apply_to_collection` is deprecated and has no effect:" - " frozen dataclasses are now reconstructed automatically. It will be removed in a future release.", - DeprecationWarning, - stacklevel=2, - ) - 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( diff --git a/tests/unittests/core/test_apply_func.py b/tests/unittests/core/test_apply_func.py index 8264ea2e..e7578f99 100644 --- a/tests/unittests/core/test_apply_func.py +++ b/tests/unittests/core/test_apply_func.py @@ -356,19 +356,6 @@ class Foo: assert foo == Foo(0) -def test_apply_to_collection_allow_frozen_kwarg_is_ignored(): - @dataclasses.dataclass(frozen=True) - class Foo: - var: int - - foo = Foo(0) - # `allow_frozen` is deprecated; passing it must warn, must not error, and must not be forwarded to `function` - with pytest.warns(DeprecationWarning, match="allow_frozen"): - result = apply_to_collection(foo, int, lambda x: x + 1, allow_frozen=True) - assert result == Foo(1) - assert foo == Foo(0) - - def test_apply_to_collection_frozen_dataclass_nested(): @dataclasses.dataclass(frozen=True) class Inner: From 06c7fad0abc7ed2d2664592c18ddf758515a3bfc Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 28 Jun 2026 09:56:13 +0000 Subject: [PATCH 5/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/lightning_utilities/core/apply_func.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/lightning_utilities/core/apply_func.py b/src/lightning_utilities/core/apply_func.py index 1eea29a6..00c370d6 100644 --- a/src/lightning_utilities/core/apply_func.py +++ b/src/lightning_utilities/core/apply_func.py @@ -3,7 +3,6 @@ # http://www.apache.org/licenses/LICENSE-2.0 # import dataclasses -import warnings from collections import OrderedDict, defaultdict from collections.abc import Callable, Mapping, Sequence from copy import deepcopy From 8e0daa31193c1d9dfd5dc39cd3c29c0b170b0aba Mon Sep 17 00:00:00 2001 From: deependujha Date: Sun, 28 Jun 2026 15:28:15 +0530 Subject: [PATCH 6/7] update --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e81f1aef..ce29fc20 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,7 +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_collections` support for frozen dataclasses with `InitVar` ([#499](https://github.com/Lightning-AI/utilities/pull/499)) +- Added `apply_to_collection`(s) support for frozen dataclasses ([#499](https://github.com/Lightning-AI/utilities/pull/499)) ### Changed From e415ceb0a018a16d433c6625846772068a815369 Mon Sep 17 00:00:00 2001 From: deependujha Date: Sun, 28 Jun 2026 15:34:03 +0530 Subject: [PATCH 7/7] update --- src/lightning_utilities/core/apply_func.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/src/lightning_utilities/core/apply_func.py b/src/lightning_utilities/core/apply_func.py index 00c370d6..a851bb62 100644 --- a/src/lightning_utilities/core/apply_func.py +++ b/src/lightning_utilities/core/apply_func.py @@ -40,7 +40,7 @@ def _dataclass_has_init_vars(data: Any) -> bool: fields. """ - return any(f._field_type is dataclasses._FIELD_INITVAR for f in data.__dataclass_fields__.values()) + 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: @@ -197,7 +197,8 @@ def _apply_field(field_value: Any) -> Any: # 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 @@ -310,7 +311,13 @@ 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.") @@ -350,9 +357,10 @@ def _apply_field(field_value1: Any) -> Any: 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(