Fix NameError in v3 from_zarr/ArraySpec.like and KeyError in model_dump

changes/169.bugfix.md

@@ -0,0 +1,5 @@
+Fixed `NameError` raised by `ArraySpec.like` and the module-level `from_zarr` in
+`pydantic_zarr.v3` and `pydantic_zarr.experimental.v3` when `zarr` was not explicitly
+imported before calling these functions (zarr was only imported under `TYPE_CHECKING`).
+Also fixed a `KeyError` raised by `ArraySpec.model_dump(exclude={'dimension_names'})` in
+both modules.

src/pydantic_zarr/experimental/v3.py

@@ -42,7 +42,7 @@
     from collections.abc import Sequence
 
     import numpy.typing as npt
-    import zarr  # noqa: TC004
+    import zarr
     from zarr.abc.store import Store
     from zarr.core.array_spec import ArrayConfigParams
 
@@ -258,8 +258,8 @@ def model_dump(self, **kwargs: Any) -> dict[str, Any]:
         """
         d = super().model_dump(**kwargs)
 
-        if d["dimension_names"] is None:
-            d.pop("dimension_names")
+        if d.get("dimension_names") is None:
+            d.pop("dimension_names", None)
         return d
 
     @classmethod
@@ -1050,6 +1050,8 @@ def from_zarr(element: zarr.Array | zarr.Group, *, depth: int = -1) -> ArraySpec
     structure of the zarr group or array.
     """
 
+    import zarr
+
     if isinstance(element, zarr.Array):
         return ArraySpec.from_zarr(element)
     else:

src/pydantic_zarr/v3.py

@@ -38,7 +38,7 @@
     from collections.abc import Sequence
 
     import numpy.typing as npt
-    import zarr  # noqa: TC004
+    import zarr
     from zarr.abc.store import Store
     from zarr.core.array_spec import ArrayConfigParams
 
@@ -250,8 +250,8 @@ def model_dump(
         # TODO: use exclude_if when we require a newer version of pydantic
         d = super().model_dump(**kwargs)
 
-        if d["dimension_names"] is None:
-            d.pop("dimension_names")
+        if d.get("dimension_names") is None:
+            d.pop("dimension_names", None)
         return d
 
     @classmethod
@@ -497,10 +497,10 @@ def like(
         """
 
         other_parsed: ArraySpec
-        if isinstance(other, zarr.Array):
+        if (zarr := sys.modules.get("zarr")) and isinstance(other, zarr.Array):
             other_parsed = ArraySpec.from_zarr(other)
         else:
-            other_parsed = other
+            other_parsed = other  # type: ignore[assignment]
 
         return model_like(self, other_parsed, include=include, exclude=exclude)
 
@@ -834,6 +834,8 @@ def from_zarr(element: zarr.Array | zarr.Group, *, depth: int = -1) -> AnyArrayS
     structure of the zarr group or array.
     """
 
+    import zarr
+
     if isinstance(element, zarr.Array):
         return ArraySpec.from_zarr(element)
     else:

tests/test_pydantic_zarr/test_experimental/test_v3.py

@@ -480,6 +480,71 @@ def test_consolidated_metadata_to_from_zarr() -> None:
     assert json.loads(store["zarr.json"].to_bytes()) == json.loads(store2["zarr.json"].to_bytes())
 
 
+def _make_array_spec_exp() -> ArraySpec:
+    """Return a minimal ArraySpec (experimental) with dimension_names=None."""
+    return ArraySpec(
+        attributes={},
+        shape=(4,),
+        data_type="uint8",
+        chunk_grid={"name": "regular", "configuration": {"chunk_shape": (4,)}},
+        chunk_key_encoding={"name": "default", "configuration": {"separator": "/"}},
+        codecs=({"name": "bytes"},),
+        fill_value=0,
+    )
+
+
+def test_exp_arrayspec_like_spec_vs_spec() -> None:
+    """
+    Regression test: experimental ArraySpec.like(other_spec) must not raise NameError.
+    """
+    spec = _make_array_spec_exp()
+    assert spec.like(spec)
+
+
+def test_exp_arrayspec_like_spec_vs_zarr_array() -> None:
+    """
+    Regression test: experimental ArraySpec.like(zarr_array) must not raise NameError.
+    Previously zarr was only imported under TYPE_CHECKING so isinstance check crashed.
+    """
+    zarr = pytest.importorskip("zarr")
+    arr = zarr.create_array(store={}, shape=(4,), dtype="uint8", zarr_format=3)
+    spec = ArraySpec.from_zarr(arr)
+    assert spec.like(arr)
+
+
+def test_exp_from_zarr_array() -> None:
+    """
+    Regression test: experimental module-level from_zarr on a zarr array must not raise NameError.
+    """
+    zarr = pytest.importorskip("zarr")
+    from pydantic_zarr.experimental.v3 import from_zarr
+
+    arr = zarr.create_array(store={}, shape=(4,), dtype="uint8", zarr_format=3)
+    result = from_zarr(arr)
+    assert isinstance(result, ArraySpec)
+
+
+def test_exp_from_zarr_group() -> None:
+    """
+    Regression test: experimental module-level from_zarr on a zarr group must not raise NameError.
+    """
+    zarr = pytest.importorskip("zarr")
+    from pydantic_zarr.experimental.v3 import from_zarr
+
+    grp = zarr.open_group(store={}, mode="w", zarr_format=3)
+    result = from_zarr(grp)
+    assert isinstance(result, GroupSpec)
+
+
+def test_exp_model_dump_exclude_dimension_names() -> None:
+    """
+    Regression test: experimental model_dump(exclude={'dimension_names'}) must not raise KeyError.
+    """
+    spec = _make_array_spec_exp()
+    d = spec.model_dump(exclude={"dimension_names"})
+    assert "dimension_names" not in d
+
+
 @pytest.mark.parametrize(
     ("dtype", "expected"),
     [

tests/test_pydantic_zarr/test_v3.py

@@ -316,6 +316,74 @@ def test_v2_chunk_key_encoding() -> None:
     )
 
 
+def _make_array_spec() -> AnyArraySpec:
+    """Return a minimal ArraySpec with dimension_names=None for regression tests."""
+    return ArraySpec(
+        attributes={},
+        shape=(4,),
+        data_type="uint8",
+        chunk_grid={"name": "regular", "configuration": {"chunk_shape": (4,)}},
+        chunk_key_encoding={"name": "default", "configuration": {"separator": "/"}},
+        codecs=({"name": "bytes"},),
+        fill_value=0,
+    )
+
+
+def test_arrayspec_like_spec_vs_spec() -> None:
+    """
+    Regression test: ArraySpec.like(other_spec) must not raise NameError.
+    Previously crashed because `zarr` was only imported under TYPE_CHECKING.
+    """
+    spec = _make_array_spec()
+    assert spec.like(spec)
+
+
+def test_arrayspec_like_spec_vs_zarr_array() -> None:
+    """
+    Regression test: ArraySpec.like(zarr_array) must not raise NameError.
+    Previously zarr was only imported under TYPE_CHECKING so isinstance check crashed.
+    """
+    zarr = pytest.importorskip("zarr")
+    arr = zarr.create_array(store={}, shape=(4,), dtype="uint8", zarr_format=3)
+    spec = ArraySpec.from_zarr(arr)
+    assert spec.like(arr)
+
+
+def test_from_zarr_array() -> None:
+    """
+    Regression test: module-level from_zarr on a zarr array must not raise NameError.
+    Previously the function body referenced `zarr.Array` without a runtime import.
+    """
+    zarr = pytest.importorskip("zarr")
+    from pydantic_zarr.v3 import from_zarr
+
+    arr = zarr.create_array(store={}, shape=(4,), dtype="uint8", zarr_format=3)
+    result = from_zarr(arr)
+    assert isinstance(result, ArraySpec)
+
+
+def test_from_zarr_group() -> None:
+    """
+    Regression test: module-level from_zarr on a zarr group must not raise NameError.
+    """
+    zarr = pytest.importorskip("zarr")
+    from pydantic_zarr.v3 import from_zarr
+
+    grp = zarr.open_group(store={}, mode="w", zarr_format=3)
+    result = from_zarr(grp)
+    assert isinstance(result, GroupSpec)
+
+
+def test_model_dump_exclude_dimension_names() -> None:
+    """
+    Regression test: model_dump(exclude={'dimension_names'}) must not raise KeyError.
+    Previously the override did d["dimension_names"] unconditionally.
+    """
+    spec = _make_array_spec()
+    d = spec.model_dump(exclude={"dimension_names"})
+    assert "dimension_names" not in d
+
+
 @pytest.mark.parametrize(
     ("dtype", "expected"),
     [