From d8a0a5533c8a31efd3ce914bbf798f6293a1e279 Mon Sep 17 00:00:00 2001
From: Norman Rzepka <code@normanrz.com>
Date: Mon, 2 Dec 2024 13:38:34 +0100
Subject: [PATCH] Adds `astype` config for Delta in Zarr 3 wrapper (#664)

* fixes #663

* changelog
---
 docs/release.rst              | 13 +++++++++++++
 docs/zarr3.rst                |  2 ++
 numcodecs/tests/test_zarr3.py | 21 +++++++++++++++++++++
 numcodecs/zarr3.py            | 14 +++++++++++++-
 4 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/docs/release.rst b/docs/release.rst
index 001f6506..738f24ec 100644
--- a/docs/release.rst
+++ b/docs/release.rst
@@ -13,6 +13,19 @@ Release notes
 
 .. _unreleased:
 
+Unreleased
+----------
+
+Fixes
+~~~~~
+* Fixes issue with ``Delta`` Zarr 3 codec not working with ``astype``.
+  By :user:`Norman Rzepka <normanrz>`, :issue:`664`
+
+
+Improvements
+~~~~~~~~~~~~
+
+
 0.14.1
 ------
 
diff --git a/docs/zarr3.rst b/docs/zarr3.rst
index d0d8c486..0596d4fd 100644
--- a/docs/zarr3.rst
+++ b/docs/zarr3.rst
@@ -1,3 +1,5 @@
+.. _Zarr 3 codecs:
+
 Zarr 3 codecs
 =============
 .. automodule:: numcodecs.zarr3
diff --git a/numcodecs/tests/test_zarr3.py b/numcodecs/tests/test_zarr3.py
index a3725296..0364dc50 100644
--- a/numcodecs/tests/test_zarr3.py
+++ b/numcodecs/tests/test_zarr3.py
@@ -244,3 +244,24 @@ def test_generic_bytes_codec(store: StorePath, codec_class: type[numcodecs.zarr3
 
     a[:, :] = data.copy()
     np.testing.assert_array_equal(data, a[:, :])
+
+
+def test_delta_astype(store: StorePath):
+    data = np.linspace(0, 10, 256, dtype="i8").reshape((16, 16))
+
+    with pytest.warns(UserWarning, match=EXPECTED_WARNING_STR):
+        a = Array.create(
+            store / "generic",
+            shape=data.shape,
+            chunk_shape=(16, 16),
+            dtype=data.dtype,
+            fill_value=0,
+            codecs=[
+                numcodecs.zarr3.Delta(dtype="i8", astype="i2"),  # type: ignore[arg-type]
+                BytesCodec(),
+            ],
+        )
+
+        a[:, :] = data.copy()
+        a = Array.open(store / "generic")
+    np.testing.assert_array_equal(data, a[:, :])
diff --git a/numcodecs/zarr3.py b/numcodecs/zarr3.py
index c4c45f1c..0fc0acee 100644
--- a/numcodecs/zarr3.py
+++ b/numcodecs/zarr3.py
@@ -266,7 +266,19 @@ def evolve_from_array_spec(self, array_spec: ArraySpec) -> Shuffle:
 
 
 # array-to-array codecs ("filters")
-Delta = _add_docstring(_make_array_array_codec("delta", "Delta"), "numcodecs.delta.Delta")
+@_add_docstring_wrapper("numcodecs.delta.Delta")
+class Delta(_NumcodecsArrayArrayCodec):
+    codec_name = f"{CODEC_PREFIX}delta"
+
+    def __init__(self, **codec_config: dict[str, JSON]) -> None:
+        super().__init__(**codec_config)
+
+    def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec:
+        if astype := self.codec_config.get("astype"):
+            return replace(chunk_spec, dtype=np.dtype(astype))  # type: ignore[arg-type]
+        return chunk_spec
+
+
 BitRound = _add_docstring(
     _make_array_array_codec("bitround", "BitRound"), "numcodecs.bitround.BitRound"
 )