From cd1a652e0852f1d6aa3f0b4f338589527eaa017e Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 9 Dec 2025 16:55:20 +0100 Subject: [PATCH 1/2] add with_* methods --- src/pydantic_zarr/experimental/v2.py | 180 +++++++++++++++++ src/pydantic_zarr/experimental/v3.py | 182 +++++++++++++++++- .../test_experimental/test_v2.py | 96 +++++++++ .../test_experimental/test_v3.py | 102 ++++++++++ 4 files changed, 559 insertions(+), 1 deletion(-) diff --git a/src/pydantic_zarr/experimental/v2.py b/src/pydantic_zarr/experimental/v2.py index a6ef1e2..cb2fd83 100644 --- a/src/pydantic_zarr/experimental/v2.py +++ b/src/pydantic_zarr/experimental/v2.py @@ -448,6 +448,150 @@ def like( return model_like(self, other_parsed, include=include, exclude=exclude) + def with_attributes(self, attributes: BaseAttributes) -> Self: + """ + Return a copy of this model with a new `attributes` field. + + Parameters + ---------- + attributes : BaseAttributes + The new `attributes` field for the copy of this model. + + Returns + ------- + ArraySpec + A copy of this model with a new `attributes` field. + """ + return type(self)(**{**self.model_dump(), "attributes": attributes}) + + def with_shape(self, shape: tuple[int, ...]) -> Self: + """ + Return a copy of this model with a new `shape` field. + + Parameters + ---------- + shape : tuple[int, ...] + The new `shape` field for the copy of this model. + + Returns + ------- + ArraySpec + A copy of this model with a new `shape` field. + """ + return type(self)(**{**self.model_dump(), "shape": shape}) + + def with_chunks(self, chunks: tuple[int, ...]) -> Self: + """ + Return a copy of this model with a new `chunks` field. + + Parameters + ---------- + chunks : tuple[int, ...] + The new `chunks` field for the copy of this model. + + Returns + ------- + ArraySpec + A copy of this model with a new `chunks` field. + """ + return type(self)(**{**self.model_dump(), "chunks": chunks}) + + def with_dtype(self, dtype: DtypeStr | list[tuple[Any, ...]]) -> Self: + """ + Return a copy of this model with a new `dtype` field. + + Parameters + ---------- + dtype : DtypeStr | list[tuple[Any, ...]] + The new `dtype` field for the copy of this model. + + Returns + ------- + ArraySpec + A copy of this model with a new `dtype` field. + """ + return type(self)(**{**self.model_dump(), "dtype": dtype}) + + def with_fill_value(self, fill_value: FillValue) -> Self: + """ + Return a copy of this model with a new `fill_value` field. + + Parameters + ---------- + fill_value : FillValue + The new `fill_value` field for the copy of this model. + + Returns + ------- + ArraySpec + A copy of this model with a new `fill_value` field. + """ + return type(self)(**{**self.model_dump(), "fill_value": fill_value}) + + def with_order(self, order: MemoryOrder) -> Self: + """ + Return a copy of this model with a new `order` field. + + Parameters + ---------- + order : MemoryOrder + The new `order` field for the copy of this model. + + Returns + ------- + ArraySpec + A copy of this model with a new `order` field. + """ + return type(self)(**{**self.model_dump(), "order": order}) + + def with_filters(self, filters: tuple[CodecDict, ...] | None) -> Self: + """ + Return a copy of this model with a new `filters` field. + + Parameters + ---------- + filters : tuple[CodecDict, ...] | None + The new `filters` field for the copy of this model. + + Returns + ------- + ArraySpec + A copy of this model with a new `filters` field. + """ + return type(self)(**{**self.model_dump(), "filters": filters}) + + def with_dimension_separator(self, dimension_separator: DimensionSeparator) -> Self: + """ + Return a copy of this model with a new `dimension_separator` field. + + Parameters + ---------- + dimension_separator : DimensionSeparator + The new `dimension_separator` field for the copy of this model. + + Returns + ------- + ArraySpec + A copy of this model with a new `dimension_separator` field. + """ + return type(self)(**{**self.model_dump(), "dimension_separator": dimension_separator}) + + def with_compressor(self, compressor: CodecDict | None) -> Self: + """ + Return a copy of this model with a new `compressor` field. + + Parameters + ---------- + compressor : CodecDict | None + The new `compressor` field for the copy of this model. + + Returns + ------- + ArraySpec + A copy of this model with a new `compressor` field. + """ + return type(self)(**{**self.model_dump(), "compressor": compressor}) + class BaseGroupSpec(StrictBase): """ @@ -457,6 +601,24 @@ class BaseGroupSpec(StrictBase): zarr_format: Literal[2] = 2 attributes: BaseAttributes + def with_attributes(self, attributes: BaseAttributes) -> Self: + """ + Return a copy of this model with a new `attributes` field. + + The new model will be validated. + + Parameters + ---------- + attributes : BaseAttributes + The new `attributes` field for the copy of this model. + + Returns + ------- + BaseGroupSpec + A copy of this model with a new `attributes` field. + """ + return type(self)(**{**self.model_dump(), "attributes": attributes}) + class GroupSpec(BaseGroupSpec): """ @@ -482,6 +644,24 @@ class can be found in the def validate_members(cls, v: BaseMember) -> BaseMember: return ensure_key_no_path(v) + def with_members(self, members: BaseMember) -> Self: + """ + Return a copy of this model with a new `members` field. + + The new model will be validated. + + Parameters + ---------- + members : Mapping[str, ArraySpec | GroupSpec] + The new `members` field for the copy of this model. + + Returns + ------- + GroupSpec + A copy of this model with a new `members` field. + """ + return type(self)(**{**self.model_dump(), "members": members}) + @classmethod def from_zarr(cls, group: zarr.Group, *, depth: int = -1) -> Self: """ diff --git a/src/pydantic_zarr/experimental/v3.py b/src/pydantic_zarr/experimental/v3.py index b1da7e4..c24b914 100644 --- a/src/pydantic_zarr/experimental/v3.py +++ b/src/pydantic_zarr/experimental/v3.py @@ -467,6 +467,150 @@ def like( return model_like(self, other_parsed, include=include, exclude=exclude) + def with_attributes(self, attributes: BaseAttributes) -> Self: + """ + Return a copy of this model with a new `attributes` field. + + Parameters + ---------- + attributes : BaseAttributes + The new `attributes` field for the copy of this model. + + Returns + ------- + ArraySpec + A copy of this model with a new `attributes` field. + """ + return type(self)(**{**self.model_dump(), "attributes": attributes}) + + def with_shape(self, shape: tuple[int, ...]) -> Self: + """ + Return a copy of this model with a new `shape` field. + + Parameters + ---------- + shape : tuple[int, ...] + The new `shape` field for the copy of this model. + + Returns + ------- + ArraySpec + A copy of this model with a new `shape` field. + """ + return type(self)(**{**self.model_dump(), "shape": shape}) + + def with_data_type(self, data_type: DTypeLike) -> Self: + """ + Return a copy of this model with a new `data_type` field. + + Parameters + ---------- + data_type : DTypeLike + The new `data_type` field for the copy of this model. + + Returns + ------- + ArraySpec + A copy of this model with a new `data_type` field. + """ + return type(self)(**{**self.model_dump(), "data_type": data_type}) + + def with_chunk_grid(self, chunk_grid: RegularChunking) -> Self: + """ + Return a copy of this model with a new `chunk_grid` field. + + Parameters + ---------- + chunk_grid : RegularChunking + The new `chunk_grid` field for the copy of this model. + + Returns + ------- + ArraySpec + A copy of this model with a new `chunk_grid` field. + """ + return type(self)(**{**self.model_dump(), "chunk_grid": chunk_grid}) + + def with_chunk_key_encoding(self, chunk_key_encoding: DefaultChunkKeyEncoding) -> Self: + """ + Return a copy of this model with a new `chunk_key_encoding` field. + + Parameters + ---------- + chunk_key_encoding : DefaultChunkKeyEncoding + The new `chunk_key_encoding` field for the copy of this model. + + Returns + ------- + ArraySpec + A copy of this model with a new `chunk_key_encoding` field. + """ + return type(self)(**{**self.model_dump(), "chunk_key_encoding": chunk_key_encoding}) + + def with_fill_value(self, fill_value: FillValue) -> Self: + """ + Return a copy of this model with a new `fill_value` field. + + Parameters + ---------- + fill_value : FillValue + The new `fill_value` field for the copy of this model. + + Returns + ------- + ArraySpec + A copy of this model with a new `fill_value` field. + """ + return type(self)(**{**self.model_dump(), "fill_value": fill_value}) + + def with_codecs(self, codecs: CodecTuple) -> Self: + """ + Return a copy of this model with a new `codecs` field. + + Parameters + ---------- + codecs : CodecTuple + The new `codecs` field for the copy of this model. + + Returns + ------- + ArraySpec + A copy of this model with a new `codecs` field. + """ + return type(self)(**{**self.model_dump(), "codecs": codecs}) + + def with_storage_transformers(self, storage_transformers: tuple[AnyNamedConfig, ...]) -> Self: + """ + Return a copy of this model with a new `storage_transformers` field. + + Parameters + ---------- + storage_transformers : tuple[AnyNamedConfig, ...] + The new `storage_transformers` field for the copy of this model. + + Returns + ------- + ArraySpec + A copy of this model with a new `storage_transformers` field. + """ + return type(self)(**{**self.model_dump(), "storage_transformers": storage_transformers}) + + def with_dimension_names(self, dimension_names: tuple[str | None, ...] | None) -> Self: + """ + Return a copy of this model with a new `dimension_names` field. + + Parameters + ---------- + dimension_names : tuple[str | None, ...] | None + The new `dimension_names` field for the copy of this model. + + Returns + ------- + ArraySpec + A copy of this model with a new `dimension_names` field. + """ + return type(self)(**{**self.model_dump(), "dimension_names": dimension_names}) + class BaseGroupSpec(StrictBase): """ @@ -476,6 +620,25 @@ class BaseGroupSpec(StrictBase): zarr_format: Literal[3] = 3 attributes: BaseAttributes + def with_attributes(self, attributes: BaseAttributes) -> Self: + """ + Return a copy of this model with a new `attributes` field. + + The new model will be validated. + + Parameters + ---------- + attributes : BaseAttributes + The new `attributes` field for the copy of this model. + + Returns + ------- + ArraySpec + A copy of this model with a new `attributes` field. + + """ + return type(self)(**{**self.model_dump(), "attributes": attributes}) + class GroupSpec(BaseGroupSpec): """ @@ -502,6 +665,23 @@ class GroupSpec(BaseGroupSpec): def validate_members(cls, v: BaseMember) -> BaseMember: return ensure_key_no_path(v) + def with_members(self, members: BaseMember) -> Self: + """ + Return a copy of this model with a new `members` field. + + The new model will be validated. + + Parameters + ---------- + members : Mapping[str, ArraySpec | GroupSpec] + The new `members` field for the copy of this model. + + Returns + ------- + A copy of this model with a new `members` field. + """ + return type(self)(**{**self.model_dump(), "members": members}) + @classmethod def from_flat(cls, data: Mapping[str, ArraySpec | BaseGroupSpec]) -> Self: """ @@ -526,7 +706,7 @@ def from_flat(cls, data: Mapping[str, ArraySpec | BaseGroupSpec]) -> Self: ```py from pydantic_zarr.experimental.v3 import GroupSpec, ArraySpec, BaseGroupSpec import numpy as np - flat = {'': BaseGroupSpec(attributes={'foo': 10})} + flat = {'' : BaseGroupSpec(attributes={'foo': 10})} GroupSpec.from_flat(flat) # GroupSpec(zarr_format=3, node_type='group', attributes={'foo': 10}, members={}) flat = { diff --git a/tests/test_pydantic_zarr/test_experimental/test_v2.py b/tests/test_pydantic_zarr/test_experimental/test_v2.py index 889966e..6a47139 100644 --- a/tests/test_pydantic_zarr/test_experimental/test_v2.py +++ b/tests/test_pydantic_zarr/test_experimental/test_v2.py @@ -596,3 +596,99 @@ def test_mix_v3_v2_fails() -> None: ), ): GroupSpec.from_flat(members_flat) # type: ignore[arg-type] + + +def test_arrayspec_with_methods() -> None: + """ + Test that ArraySpec with_* methods create new validated copies + """ + original = ArraySpec.from_array(np.arange(10), attributes={"foo": "bar"}) + + # Test with_attributes + new_attrs = original.with_attributes({"baz": "qux"}) + assert new_attrs.attributes == {"baz": "qux"} + assert original.attributes == {"foo": "bar"} # Original unchanged + assert new_attrs is not original + + # Test with_shape + new_shape = original.with_shape((20,)) + assert new_shape.shape == (20,) + assert original.shape == (10,) + + # Test with_chunks + new_chunks = original.with_chunks((5,)) + assert new_chunks.chunks == (5,) + assert original.chunks == (10,) + + # Test with_dtype + new_dtype = original.with_dtype("float32") + assert new_dtype.dtype == " None: + """ + Test that ArraySpec with_* methods trigger validation + """ + spec = ArraySpec(shape=(10,), chunks=(5,), dtype="uint8", attributes={}) + + # Test that validation fails when shape and chunks have mismatched lengths + with pytest.raises(ValidationError): + spec.with_shape((10, 10)) # Shape has 2 dims but chunks still has 1 + + +def test_groupspec_with_methods() -> None: + """ + Test that GroupSpec with_* methods create new validated copies + """ + array_spec = ArraySpec.from_array(np.arange(10), attributes={}) + original = GroupSpec(attributes={"group": "attr"}, members={"arr": array_spec}) + + # Test with_attributes + new_attrs = original.with_attributes({"new": "attr"}) + assert new_attrs.attributes == {"new": "attr"} + assert original.attributes == {"group": "attr"} # Original unchanged + assert new_attrs is not original + + # Test with_members + new_array = ArraySpec.from_array(np.arange(5), attributes={}) + new_members = original.with_members({"new_arr": new_array}) + assert "new_arr" in new_members.members + assert "arr" not in new_members.members # Replacement, not merge + assert "arr" in original.members # Original unchanged + + +def test_groupspec_with_members_validation() -> None: + """ + Test that GroupSpec with_members triggers validation + """ + spec = GroupSpec(attributes={}, members={}) + + # Test that validation fails with invalid member names + with pytest.raises(ValidationError, match='Strings containing "/" are invalid'): + spec.with_members({"a/b": ArraySpec.from_array(np.arange(10), attributes={})}) diff --git a/tests/test_pydantic_zarr/test_experimental/test_v3.py b/tests/test_pydantic_zarr/test_experimental/test_v3.py index 3376419..2e1bf4e 100644 --- a/tests/test_pydantic_zarr/test_experimental/test_v3.py +++ b/tests/test_pydantic_zarr/test_experimental/test_v3.py @@ -302,3 +302,105 @@ def test_dim_names_from_zarr_array( arr = zarr.zeros(*args, **kwargs) spec: ArraySpec = ArraySpec.from_zarr(arr) assert spec.dimension_names == expected_names + + +def test_arrayspec_with_methods() -> None: + """ + Test that ArraySpec with_* methods create new validated copies + """ + original = ArraySpec.from_array(np.arange(10), attributes={"foo": "bar"}) + + # Test with_attributes + new_attrs = original.with_attributes({"baz": "qux"}) + assert new_attrs.attributes == {"baz": "qux"} + assert original.attributes == {"foo": "bar"} # Original unchanged + assert new_attrs is not original + + # Test with_shape + new_shape = original.with_shape((20,)) + assert new_shape.shape == (20,) + assert original.shape == (10,) + + # Test with_data_type + new_dtype = original.with_data_type("float32") + assert new_dtype.data_type == "float32" + assert original.data_type == "int64" + + # Test with_chunk_grid + new_grid = original.with_chunk_grid({"name": "regular", "configuration": {"chunk_shape": (5,)}}) + assert new_grid.chunk_grid["configuration"]["chunk_shape"] == (5,) # type: ignore[index] + assert original.chunk_grid["configuration"]["chunk_shape"] == (10,) # type: ignore[index] + + # Test with_chunk_key_encoding + new_encoding = original.with_chunk_key_encoding( + {"name": "default", "configuration": {"separator": "."}} + ) + assert new_encoding.chunk_key_encoding["configuration"]["separator"] == "." # type: ignore[index] + assert original.chunk_key_encoding["configuration"]["separator"] == "/" # type: ignore[index] + + # Test with_fill_value + new_fill = original.with_fill_value(999) + assert new_fill.fill_value == 999 + assert original.fill_value == 0 + + # Test with_codecs + new_codecs = original.with_codecs(({"name": "gzip", "configuration": {"level": 1}},)) + assert len(new_codecs.codecs) == 1 + assert new_codecs.codecs[0]["name"] == "gzip" # type: ignore[index] + + # Test with_storage_transformers + new_transformers = original.with_storage_transformers(({"name": "test", "configuration": {}},)) + assert len(new_transformers.storage_transformers) == 1 + assert original.storage_transformers == () + + # Test with_dimension_names + new_dims = original.with_dimension_names(("x",)) + assert new_dims.dimension_names == ("x",) + assert original.dimension_names is None + + +def test_arrayspec_with_methods_validation() -> None: + """ + Test that ArraySpec with_* methods trigger validation + """ + spec = ArraySpec.from_array(np.arange(10), attributes={}) + + # Test that validation fails when dimension_names length doesn't match shape + with pytest.raises(ValidationError, match="Invalid `dimension names` attribute"): + spec.with_dimension_names(("x", "y")) # 2 names for 1D array + + # Test that validation fails with empty codecs + with pytest.raises(ValidationError, match="Invalid length. Expected 1 or more, got 0"): + spec.with_codecs(()) + + +def test_groupspec_with_methods() -> None: + """ + Test that GroupSpec with_* methods create new validated copies + """ + array_spec = ArraySpec.from_array(np.arange(10), attributes={}) + original = GroupSpec(attributes={"group": "attr"}, members={"arr": array_spec}) + + # Test with_attributes + new_attrs = original.with_attributes({"new": "attr"}) + assert new_attrs.attributes == {"new": "attr"} + assert original.attributes == {"group": "attr"} # Original unchanged + assert new_attrs is not original + + # Test with_members + new_array = ArraySpec.from_array(np.arange(5), attributes={}) + new_members = original.with_members({"new_arr": new_array}) + assert "new_arr" in new_members.members + assert "arr" not in new_members.members # Replacement, not merge + assert "arr" in original.members # Original unchanged + + +def test_groupspec_with_members_validation() -> None: + """ + Test that GroupSpec with_members triggers validation + """ + spec = GroupSpec(attributes={}, members={}) + + # Test that validation fails with invalid member names + with pytest.raises(ValidationError, match='Strings containing "/" are invalid'): + spec.with_members({"a/b": ArraySpec.from_array(np.arange(10), attributes={})}) From 60f0f11efb01ec28d8bba822b12e4a57e484c9c6 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 9 Dec 2025 16:57:47 +0100 Subject: [PATCH 2/2] better type annotations --- src/pydantic_zarr/experimental/v3.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/pydantic_zarr/experimental/v3.py b/src/pydantic_zarr/experimental/v3.py index c24b914..7ab9d04 100644 --- a/src/pydantic_zarr/experimental/v3.py +++ b/src/pydantic_zarr/experimental/v3.py @@ -515,7 +515,7 @@ def with_data_type(self, data_type: DTypeLike) -> Self: """ return type(self)(**{**self.model_dump(), "data_type": data_type}) - def with_chunk_grid(self, chunk_grid: RegularChunking) -> Self: + def with_chunk_grid(self, chunk_grid: AnyNamedConfig) -> Self: """ Return a copy of this model with a new `chunk_grid` field. @@ -531,7 +531,7 @@ def with_chunk_grid(self, chunk_grid: RegularChunking) -> Self: """ return type(self)(**{**self.model_dump(), "chunk_grid": chunk_grid}) - def with_chunk_key_encoding(self, chunk_key_encoding: DefaultChunkKeyEncoding) -> Self: + def with_chunk_key_encoding(self, chunk_key_encoding: AnyNamedConfig) -> Self: """ Return a copy of this model with a new `chunk_key_encoding` field.