make examples stand-alone and testable via script dependency modification at test time

Author: d-v-b

examples/custom_dtype.py

@@ -1,7 +1,7 @@
 # /// script
 # requires-python = ">=3.11"
 # dependencies = [
-#   "zarr @ {root}",
+#   "zarr @ git+https://github.com/zarr-developers/zarr-python.git@main",
 #   "ml_dtypes==0.5.1",
 #   "pytest==8.4.1"
 # ]
@@ -18,7 +18,7 @@
 import json
 import sys
 from pathlib import Path
-from typing import ClassVar, Literal, Self, TypeGuard
+from typing import ClassVar, Literal, Self, TypeGuard, overload
 
 import ml_dtypes  # necessary to add extra dtypes to NumPy
 import numpy as np
@@ -34,14 +34,17 @@
     check_dtype_spec_v2,
 )
 
+# This is the int2 array data type
 int2_dtype_cls = type(np.dtype("int2"))
+
+# This is the int2 scalar type
 int2_scalar_cls = ml_dtypes.int2
 
 
 class Int2(ZDType[int2_dtype_cls, int2_scalar_cls]):
     """
-    This class provides a Zarr compatibility layer around the int2 data type and the int2
-    scalar type.
+    This class provides a Zarr compatibility layer around the int2 data type (the ``dtype`` of a
+    NumPy array of type int2) and the int2 scalar type (the ``dtype`` of the scalar value inside an int2 array).
     """
 
     # This field is as the key for the data type in the internal data type registry, and also
@@ -68,72 +71,140 @@ def to_native_dtype(self: Self) -> int2_dtype_cls:
 
     @classmethod
     def _check_json_v2(cls, data: DTypeJSON) -> TypeGuard[DTypeConfig_V2[Literal["|b1"], None]]:
-        """Type check for Zarr v2-flavored JSON"""
+        """
+        Type check for Zarr v2-flavored JSON.
+
+        This will check that the input is a dict like this:
+        .. code-block:: json
+
+        {
+            "name": "int2",
+            "object_codec_id": None
+        }
+
+        Note that this representation differs from the ``dtype`` field looks like in zarr v2 metadata.
+        Specifically, whatever goes into the ``dtype`` field in metadata is assigned to the ``name`` field here.
+
+        See the Zarr docs for more information about the JSON encoding for data types.
+        """
         return (
             check_dtype_spec_v2(data) and data["name"] == "int2" and data["object_codec_id"] is None
         )
 
     @classmethod
     def _check_json_v3(cls, data: DTypeJSON) -> TypeGuard[Literal["int2"]]:
-        """Type check for Zarr v3-flavored JSON"""
+        """
+        Type check for Zarr V3-flavored JSON.
+
+        Checks that the input is the string "int2".
+        """
         return data == cls._zarr_v3_name
 
     @classmethod
     def _from_json_v2(cls, data: DTypeJSON) -> Self:
         """
-        Create an instance of this ZDType from zarr v3-flavored JSON.
+        Create an instance of this ZDType from Zarr V3-flavored JSON.
         """
         if cls._check_json_v2(data):
             return cls()
+        #  This first does a type check on the input, and if that passes we create an instance of the ZDType.
         msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected the string {cls._zarr_v2_name!r}"
         raise DataTypeValidationError(msg)
 
     @classmethod
     def _from_json_v3(cls: type[Self], data: DTypeJSON) -> Self:
         """
-        Create an instance of this ZDType from zarr v3-flavored JSON.
+        Create an instance of this ZDType from Zarr V3-flavored JSON.
+
+        This first does a type check on the input, and if that passes we create an instance of the ZDType.
         """
         if cls._check_json_v3(data):
             return cls()
         msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected the string {cls._zarr_v3_name!r}"
         raise DataTypeValidationError(msg)
 
+    @overload  # type: ignore[override]
+    def to_json(self, zarr_format: Literal[2]) -> DTypeConfig_V2[Literal["int2"], None]: ...
+
+    @overload
+    def to_json(self, zarr_format: Literal[3]) -> Literal["int2"]: ...
+
     def to_json(
         self, zarr_format: ZarrFormat
     ) -> DTypeConfig_V2[Literal["int2"], None] | Literal["int2"]:
-        """Serialize this ZDType to v2- or v3-flavored JSON"""
+        """
+        Serialize this ZDType to v2- or v3-flavored JSON
+
+        If the zarr_format is 2, then return a dict like this:
+        .. code-block:: json
+
+        {
+            "name": "int2",
+            "object_codec_id": None
+        }
+
+        If the zarr_format is 3, then return the string "int2"
+
+        """
         if zarr_format == 2:
             return {"name": "int2", "object_codec_id": None}
         elif zarr_format == 3:
             return self._zarr_v3_name
         raise ValueError(f"zarr_format must be 2 or 3, got {zarr_format}")  # pragma: no cover
 
-    def _check_scalar(self, data: object) -> TypeGuard[int]:
-        """Check if a python object is a valid scalar"""
+    def _check_scalar(self, data: object) -> TypeGuard[int | ml_dtypes.int2]:
+        """
+        Check if a python object is a valid int2-compatible scalar
+
+        The strictness of this type check is an implementation degree of freedom.
+        You could be strict here, and only accept int2 values, or be open and accept any integer
+        or any object and rely on exceptions from the int2 constructor that will be called in
+        cast_scalar.
+        """
         return isinstance(data, (int, int2_scalar_cls))
 
     def cast_scalar(self, data: object) -> ml_dtypes.int2:
         """
-        Attempt to cast a python object to an int2. Might fail pending a type check.
+        Attempt to cast a python object to an int2.
+
+        We first perform a type check to ensure that the input type is appropriate, and if that
+        passes we call the int2 scalar constructor.
         """
         if self._check_scalar(data):
             return ml_dtypes.int2(data)
         msg = f"Cannot convert object with type {type(data)} to a 2-bit integer."
         raise TypeError(msg)
 
     def default_scalar(self) -> ml_dtypes.int2:
-        """Get the default scalar value"""
+        """
+        Get the default scalar value. This will be used when automatically selecting a fill value.
+        """
         return ml_dtypes.int2(0)
 
     def to_json_scalar(self, data: object, *, zarr_format: ZarrFormat) -> int:
-        """Convert a python object to a scalar."""
-        return int(data)
+        """
+        Convert a python object to a JSON representation of an int2 scalar.
+        This is necessary for taking user input for the ``fill_value`` attribute in array metadata.
+
+        In this implementation, we optimistically convert the input to an int,
+        and then check that it lies in the acceptable range for this data type.
+        """
+        # We could add a type check here, but we don't need to for this example
+        val: int = int(data)  # type: ignore[call-overload]
+        if val not in (-2, -1, 0, 1):
+            raise ValueError("Invalid value. Expected -2, -1, 0, or 1.")
+        return val
 
     def from_json_scalar(self, data: JSON, *, zarr_format: ZarrFormat) -> ml_dtypes.int2:
         """
-        Read a JSON-serializable value as a scalar. The base definition of this method
-        requires that it take a zarr_format parameter, because some data types serialize scalars
-        differently in zarr v2 and v3
+        Read a JSON-serializable value as an int2 scalar.
+
+        We first perform a type check to ensure that the JSON value is well-formed, then call the
+        int2 scalar constructor.
+
+        The base definition of this method requires that it take a zarr_format parameter because
+        other data types serialize scalars differently in zarr v2 and v3, but we don't use this here.
+
         """
         if self._check_scalar(data):
             return ml_dtypes.int2(data)
@@ -167,4 +238,8 @@ def test_custom_dtype(tmp_path: Path, zarr_format: Literal[2, 3]) -> None:
 
 
 if __name__ == "__main__":
+    # Run the example with printed output, and a dummy pytest configuration file specified.
+    # Without the dummy configuration file, at test time pytest will attempt to use the
+    # configuration file in the project root, which will error because Zarr is using some
+    # plugins that are not installed in this example.
     sys.exit(pytest.main(["-s", __file__, f"-c {__file__}"]))

pyproject.toml

@@ -80,6 +80,9 @@ test = [
     "mypy",
     "hypothesis",
     "pytest-xdist",
+    "packaging",
+    "tomlkit",
+    "uv"
 ]
 remote_tests = [
     'zarr[remote]',

tests/test_examples.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+import re
+import subprocess
+from pathlib import Path
+from typing import Final
+
+import pytest
+import tomlkit
+from packaging.requirements import Requirement
+
+examples_dir = "examples"
+script_paths = Path(examples_dir).glob("*.py")
+
+PEP_723_REGEX: Final = r"(?m)^# /// (?P<type>[a-zA-Z0-9-]+)$\s(?P<content>(^#(| .*)$\s)+)^# ///$"
+
+# This is the absolute path to the local Zarr installation. Moving this test to a different directory will break it.
+ZARR_PROJECT_PATH = Path(".").absolute()
+
+
+def set_dep(script: str, dependency: str) -> str:
+    """
+    Set a dependency in a PEP-723 script header.
+    If the package is already in the list, it will be replaced.
+    If the package is not already in the list, it will be added.
+
+    Source code modified from
+    https://packaging.python.org/en/latest/specifications/inline-script-metadata/#reference-implementation
+    """
+    match = re.search(PEP_723_REGEX, script)
+
+    if match is None:
+        raise ValueError(f"PEP-723 header not found in {script}")
+
+    content = "".join(
+        line[2:] if line.startswith("# ") else line[1:]
+        for line in match.group("content").splitlines(keepends=True)
+    )
+
+    config = tomlkit.parse(content)
+    for idx, dep in enumerate(tuple(config["dependencies"])):
+        if Requirement(dep).name == Requirement(dependency).name:
+            config["dependencies"][idx] = dependency
+
+    new_content = "".join(
+        f"# {line}" if line.strip() else f"#{line}"
+        for line in tomlkit.dumps(config).splitlines(keepends=True)
+    )
+
+    start, end = match.span("content")
+    return script[:start] + new_content + script[end:]
+
+
+def resave_script(source_path: Path, dest_path: Path) -> None:
+    """
+    Read a script from source_path and save it to dest_path after inserting the absolute path to the
+    local Zarr project directory in the PEP-723 header.
+    """
+    source_text = source_path.read_text()
+    dest_text = set_dep(source_text, f"zarr @ file:///{ZARR_PROJECT_PATH}")
+    dest_path.write_text(dest_text)
+
+
+@pytest.mark.parametrize("script_path", script_paths)
+def test_scripts_can_run(script_path: Path, tmp_path: Path) -> None:
+    dest_path = tmp_path / script_path.name
+    # We resave the script after inserting the absolute path to the local Zarr project directory,
+    # and then test its behavior.
+    # This allows the example to be useful to users who don't have Zarr installed, but also testable.
+    resave_script(script_path, dest_path)
+    result = subprocess.run(["uv", "run", str(dest_path)], capture_output=True, text=True)
+    assert result.returncode == 0, (
+        f"Script at {script_path} failed to run. Output: {result.stdout} Error: {result.stderr}"
+    )