From 3e98c3b1d6ed24b7b04fdbaa56ed2c42f5b1ef05 Mon Sep 17 00:00:00 2001
From: Ryan Abernathey <ryan.abernathey@gmail.com>
Date: Sun, 26 May 2024 14:54:33 -0400
Subject: [PATCH] Add zstd to old V3 supported codecs (#1914)

* add zstd to old V3 supported codecs

* get to full test coverage and add release note

* fix pre-commit
---
 docs/release.rst        |  3 +++
 zarr/codecs.py          |  2 +-
 zarr/meta.py            |  4 ++++
 zarr/tests/test_meta.py | 23 +++++++++++++++++------
 4 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/docs/release.rst b/docs/release.rst
index 59051bbf97..48a6c9ca2d 100644
--- a/docs/release.rst
+++ b/docs/release.rst
@@ -32,6 +32,9 @@ Docs
 Maintenance
 ~~~~~~~~~~~
 
+* Add Zstd codec to old V3 code path.
+  By :user:`Ryan Abernathey <rabernat>` 
+
 Deprecations
 ~~~~~~~~~~~~
 
diff --git a/zarr/codecs.py b/zarr/codecs.py
index 4ad68b8627..6fd5e20401 100644
--- a/zarr/codecs.py
+++ b/zarr/codecs.py
@@ -1,4 +1,4 @@
 # flake8: noqa
 from numcodecs import *
-from numcodecs import get_codec, Blosc, Pickle, Zlib, Delta, AsType, BZ2
+from numcodecs import get_codec, Blosc, Pickle, Zlib, Zstd, Delta, AsType, BZ2
 from numcodecs.registry import codec_registry
diff --git a/zarr/meta.py b/zarr/meta.py
index 5430ab305d..44a2b7ebec 100644
--- a/zarr/meta.py
+++ b/zarr/meta.py
@@ -414,6 +414,8 @@ def _encode_codec_metadata(cls, codec: Codec) -> Optional[Mapping]:
             uri = uri + "lz4/1.0"
         elif isinstance(codec, numcodecs.LZMA):
             uri = uri + "lzma/1.0"
+        elif isinstance(codec, numcodecs.Zstd):
+            uri = uri + "zstd/1.0"
         meta = {
             "codec": uri,
             "configuration": config,
@@ -439,6 +441,8 @@ def _decode_codec_metadata(cls, meta: Optional[Mapping]) -> Optional[Codec]:
             conf["id"] = "lz4"
         elif meta["codec"].startswith(uri + "lzma/"):
             conf["id"] = "lzma"
+        elif meta["codec"].startswith(uri + "zstd/"):
+            conf["id"] = "zstd"
         else:
             raise NotImplementedError
 
diff --git a/zarr/tests/test_meta.py b/zarr/tests/test_meta.py
index f9010d6788..7b7d526476 100644
--- a/zarr/tests/test_meta.py
+++ b/zarr/tests/test_meta.py
@@ -5,7 +5,7 @@
 import numpy as np
 import pytest
 
-from zarr.codecs import Blosc, Delta, Pickle, Zlib
+from zarr.codecs import Blosc, Delta, Pickle, Zlib, Zstd
 from zarr.errors import MetadataError
 from zarr.meta import (
     ZARR_FORMAT,
@@ -268,17 +268,23 @@ def test_encode_decode_array_dtype_shape():
     assert meta_dec["filters"] is None
 
 
-def test_encode_decode_array_dtype_shape_v3():
+@pytest.mark.parametrize("cname", ["zlib", "zstd"])
+def test_encode_decode_array_dtype_shape_v3(cname):
+    if cname == "zlib":
+        compressor = Zlib(1)
+    elif cname == "zstd":
+        compressor = Zstd(1)
     meta = dict(
         shape=(100,),
         chunk_grid=dict(type="regular", chunk_shape=(10,), separator=("/")),
         data_type=np.dtype("(10, 10)<f8"),
-        compressor=Zlib(1),
+        compressor=compressor,
         fill_value=None,
         chunk_memory_layout="C",
     )
 
-    meta_json = """{
+    meta_json = (
+        """{
         "attributes": {},
         "chunk_grid": {
             "chunk_shape": [10],
@@ -287,7 +293,11 @@ def test_encode_decode_array_dtype_shape_v3():
         },
         "chunk_memory_layout": "C",
         "compressor": {
-            "codec": "https://purl.org/zarr/spec/codec/zlib/1.0",
+        """
+        + f"""
+        "codec": "https://purl.org/zarr/spec/codec/{cname}/1.0",
+        """
+        + """
             "configuration": {
                 "level": 1
             }
@@ -297,6 +307,7 @@ def test_encode_decode_array_dtype_shape_v3():
         "fill_value": null,
         "shape": [100, 10, 10 ]
     }"""
+    )
 
     # test encoding
     meta_enc = Metadata3.encode_array_metadata(meta)
@@ -315,7 +326,7 @@ def test_encode_decode_array_dtype_shape_v3():
     assert "filters" not in meta_dec
 
 
-@pytest.mark.parametrize("comp_id", ["gzip", "zlib", "blosc", "bz2", "lz4", "lzma"])
+@pytest.mark.parametrize("comp_id", ["gzip", "zlib", "blosc", "bz2", "lz4", "lzma", "zstd"])
 def test_decode_metadata_implicit_compressor_config_v3(comp_id):
     meta = {
         "attributes": {},