add a runtime type checker for metadata objects #3400

d-v-b · 2025-08-24T13:21:51Z

This PR adds a runtime type checker specifically for checking JSON-like data against a type definition. It's currently a draft while I get the test suite happy and refine the API, but it's also ready for people to look at and try out. I'm pretty convinced of it's utility, but I also think we should have a good discussion about whether this feature is a good idea.

Demo

The basic API looks like this:

# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "zarr@git+https://github.com/d-v-b/zarr-python.git@feat/type-checker",
# ]
# ///

from zarr.core.type_check import check_type
from typing import TypedDict, Literal, NotRequired

class MyChunkGridConfig(TypedDict):
   prop_a: int
   prop_b: str
   prop_c: str | int
   optional_prop: NotRequired[bool]

class MyChunkGrid(TypedDict):
   name: Literal["awesome_chunk_grid"]
   configuration: MyChunkGridConfig

valid_data = [
   {"name": "awesome_chunk_grid", "configuration": {"prop_a": 10, "prop_b": "string", "prop_c": 10}},
   {"name": "awesome_chunk_grid", "configuration": {"prop_a": 10, "prop_b": "string", "prop_c": "10"}},
   {"name": "awesome_chunk_grid", "configuration": {"prop_a": 10, "prop_b": "string", "prop_c": "10", "optional_prop": True}}
   ]

for d in valid_data:
    result = check_type(d, MyChunkGrid)
    print(result.success, result.errors)
"""
True []
True []
True []
"""

# prop_b should be a string, but it's an int
invalid_data = [
    # invalid type
   {"name": "awesome_chunk_grid", "configuration": {"prop_a": 10, "prop_b": 10, "prop_c": 10}},
   # missing key
   {"name": "awesome_chunk_grid", "configuration": {"prop_a": 10, "prop_b": "10"}}
   ]
for d in invalid_data:
    result = check_type(d, MyChunkGrid)
    print(result.success, result.errors)
"""
False ["value['configuration']['prop_b'] expected an instance of <class 'str'> but got 10 with type <class 'int'>"]
False ["value['configuration'] missing required key 'prop_c'"]
"""

Some aspects might evolve while this is a draft, like the nature of the error messages.

Supported types

This is not a general-purpose type checker. It is targeted for the types relevant for Zarr metadata documents, and so It supports the following narrow set of types

int, bool, float, str
union
tuple
list
sequence
mapping
typeddict

cost

maintenance burden

The type checker itself is ~530 lines of commented code, broken up into functions which are mostly easy to understand. The typeddict part, and the logic for resolving generic types, is convoluted and potentially sensitive to changes in how python exposes type annotations at runtime. Many type annotation features have been designed for static type checkers and not use within a python program, so some of this is rather fiddly. But I don't think we are relying on any brittle or private APIs here.

performance

As currrently implemented, the type checker will report all detectable errors:

>>> check_type(tuple(range(4)), tuple[str, ...])
TypeCheckResult(success=False, errors=["value[0] expected an instance of <class 'str'> but got 0 with type <class 'int'>", "value[1] expected an instance of <class 'str'> but got 1 with type <class 'int'>", "value[2] expected an instance of <class 'str'> but got 2 with type <class 'int'>", "value[3] expected an instance of <class 'str'> but got 3 with type <class 'int'>"])

This is wasted compute when we don't care exactly how mismatched the data is, but it is a better user experience. We might need to tune this if performance becomes a problem, e.g. by introducing a "fail_fast" option that returns on the first error.

benefit

We can instantly remove a lot of special-purpose functions. Most of the functions named parse_* (~30+ functions) and essentially all of the functions named *check_json* (~30 functions) could be replaced or simplified with the check_type function.

We can also make our JSON loading routines type-safe:

# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "zarr@git+https://github.com/d-v-b/zarr-python.git@feat/type-checker",
# ]
# ///

from zarr.core.type_check import guard_type
from zarr.core.common import ArrayMetadataJSON_V3

unknown_data: dict[str, object] = {
    "zarr_format": 3,
    "node_type": "array",
    "shape": (10,10),
    "data_type": "uint8", 
    "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": (10,10)}},
    "chunk_key_encoding": {"name": "default", "configuration": {"separator": "/"}},
    "codecs": ({"name": "bytes"}, ),
    "fill_value": 0
    }

if guard_type(unknown_data, ArrayMetadataJSON_V3):
    reveal_type(unknown_data)

mypy could not infer the type correctly, but basedpyright does:

zarr-python git:(feat/type-checker) ✗ uvx basedpyright test.py
/Users/d-v-b/dev/zarr-python/test.py
  /Users/d-v-b/dev/zarr-python/test.py:25:17 - information: Type of "unknown_data" is "ArrayMetadataJSON_V3"
0 errors, 0 warnings, 1 note

While we could write a bespoke function that specifically checks all the possibilities for zarr v3 metadata. But then we would need to painfully modify that function by hand to support something like this:

# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "zarr@git+https://github.com/d-v-b/zarr-python.git@feat/type-checker",
# ]
# ///
from collections.abc import Mapping
from typing_extensions import TypedDict
from zarr.core.common import ArrayMetadataJSON_V3
from zarr.core.type_check import check_type

class GeoZarrAttrs(TypedDict):
    geozarr: Mapping[str, object]

class GeoZarrArray(ArrayMetadataJSON_V3):
    attributes: GeoZarrAttrs

unknown_data: dict[str, object] = {
    "zarr_format": 3,
    "node_type": "array",
    "shape": (10,10),
    "data_type": "uint8", 
    "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": (10,10)}},
    "chunk_key_encoding": {"name": "default", "configuration": {"separator": "/"}},
    "codecs": ({"name": "bytes"}, ),
    "fill_value": 0,
    "attributes": {"not_geozarr": "bar"}
    }

result = check_type(unknown_data, GeoZarrArray)
print(result)
"""
TypeCheckResult(success=False, errors=["value['attributes'] missing required key 'geozarr'"])
"""

alternatives

we could use an external JSON validation library / type checking like pydantic, attrs, msgspec, beartype, etc. But I would rather not add a dependency. With the approach in this PR, we keep control in-house, and because this PR just adds functions, it composes with the rest of our codebase at the moment. (FWIW right now this type checker doesn't do any parsing, it only validates. If you think we should parse instead of just validating, then IMO that's a job for our array metadata classes)

we could also do nothing, and continue writing JSON parsing code by hand. But I would rather not do that, because this invites bugs and makes it hard to keep up with sneaky spec changes. Specifically, I'm planning on writing a lot of new types to model the codecs defined in #3376, and I would rather just write the type and get the type checking (and type safety) for free.

closes #3285

…eddicts

…ebase

…at/type-checker

d-v-b · 2025-08-25T13:17:42Z

this is pretty substantial so I would appreciate a lot of eyes @zarr-developers/python-core-devs

if anyone has concerns about whether we should do any runtime type checking at all, maybe send those thoughts to the issue this PR closes

I'm going to keep working on tests for the type checker, but so far it's working great.

This PR does violating liskov for a few subclasses of our Metadata ABC, because that class requires that to_dict and from_dict use dict[st, JSON], which is not very accurate. After this PR we need to make changes to class so that it supports type safety, then we won't be violating liskov any more.

Similarly, there are lots of # type: ignores added in various places. As you might guess, those were added because without them, mypy flagging type errors. Many of these checks will go away when we fix the lax typing of the Metadata ABC, and other classes.

@TomAugspurger I think you in particular will appreciate some of the effects of this PR. Since we can annotate methods like ArrayV3Metadata.from_dict() as taking ArrayMetadataJSON_V3, we don't need to do any runtime validation inside from_dict. The assumption is that the caller has already checked the input. This is not possible without the types defined in this PR, and it's not practical without the type checker. Once we extrapolate the type safe style, we can push almost all our type checking to the IO boundary.

That being said, I think the ArrayMetadata class will still need to do some internal consistency checks, like ensuring that the number of dimension names matches the length of shape. I don't think we want our type checker to be smart enough to catch that kind of thing.

…valuate_forward_ref

codecov · 2025-08-25T13:37:31Z

Codecov Report

❌ Patch coverage is 98.01980% with 8 lines in your changes missing coverage. Please review.
✅ Project coverage is 94.90%. Comparing base (7d43f32) to head (c7096b1).

Files with missing lines	Patch %	Lines
src/zarr/core/type_check.py	97.95%	5 Missing ⚠️
src/zarr/core/array.py	86.66%	2 Missing ⚠️
src/zarr/core/metadata/v3.py	96.15%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3400      +/-   ##
==========================================
- Coverage   94.92%   94.90%   -0.03%     
==========================================
  Files          79       80       +1     
  Lines        9503     9752     +249     
==========================================
+ Hits         9021     9255     +234     
- Misses        482      497      +15

Files with missing lines	Coverage Δ
src/zarr/abc/codec.py	`98.50% <100.00%> (-0.22%)`	⬇️
src/zarr/api/asynchronous.py	`90.00% <100.00%> (-0.04%)`	⬇️
src/zarr/core/common.py	`95.40% <100.00%> (+2.06%)`	⬆️
src/zarr/core/dtype/common.py	`70.58% <100.00%> (-14.96%)`	⬇️
src/zarr/core/dtype/npy/bool.py	`100.00% <100.00%> (ø)`
src/zarr/core/dtype/npy/bytes.py	`99.49% <100.00%> (-0.01%)`	⬇️
src/zarr/core/dtype/npy/complex.py	`98.80% <100.00%> (+0.01%)`	⬆️
src/zarr/core/dtype/npy/float.py	`98.91% <100.00%> (+0.01%)`	⬆️
src/zarr/core/dtype/npy/int.py	`99.37% <100.00%> (+<0.01%)`	⬆️
src/zarr/core/dtype/npy/string.py	`97.79% <100.00%> (+0.01%)`	⬆️
... and 10 more

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

…at/type-checker

…iguating ints from bools

d-v-b · 2025-08-29T15:40:12Z

In the dev meeting today @jhamman rightly challenged me to justify writing our own in-house runtime type checker instead of bringing in pydantic. I had a few answers:

The type checker here is small (~500 LoC, with docstrings) and it has a very simple structure -- a bunch of functions. Pydantic is huge and complex. It relies on dynamic class creation, and bugs in in pydantic can be hard to debug.
pydantic is class-centric. I think a functional API for type checking is a better fit for our codebase. Replacing all our classes with BaseModel is relatively disruptive comparing to bringing in a few functions to help with a type guard.
Pydantic contains a lot of stuff we don't need -- For zarr, we only need to check the types that could fit in zarr.json, but pydantic covers a lot more. The scale of pydantic is also an upside -- it means that new features will get added, and bugs will be squashed, by a large number of developers. But as long as we test our in-house type checker extensively against the small set of types that matter to us, I don't think we would gain much from adding pydantic as a dependency.
Pydantic makes breaking changes. If we depend on pydantic, we then become responsible for managing how those breaking changes impact our users. This is a cost that I think we can avoid.
The changes here are very non-invasive (because it's just functions). We can always switch to pydantic later, if this type checker sucks. We could also revert and replace all instances of this type checker with a hand-written type checking function.

My calculation was, 500 LoC for an API that does exactly what we need, (and nothing else), was a better deal than bringing in a big, celebrity dependency that would require pretty substantial changes to basic Zarr Python classes. Other folks might see this calculation differently.

…at/type-checker

d-v-b · 2025-08-29T21:40:10Z

as illustration, consider this diff enabled by this PR:

@@ -286,14 +283,7 @@ class NullTerminatedBytes(ZDType[np.dtypes.BytesDType[int], np.bytes_], HasLengt
             True if the input is a valid representation of this class in Zarr V3, False
             otherwise.
         """
-        return (
-            isinstance(data, dict)
-            and set(data.keys()) == {"name", "configuration"}
-            and data["name"] == cls._zarr_v3_name
-            and isinstance(data["configuration"], dict)
-            and "length_bytes" in data["configuration"]
-            and isinstance(data["configuration"]["length_bytes"], int)
-        )
+        return check_type(data, NullTerminatedBytesJSON_V3).success

I worry that if type safety requires writing tedious, error-prone functions like what we have in main, we might not do it, and then we lose type safety. Boiling it down to a much simpler function makes type safety look a lot easier.

TomAugspurger

Gave a Quick Look, but haven't gone through in detail. At a high level:

I'd prefer to avoid a dependency on pydantic (or any other runtime validators)
At a glance, the implementation looks reasonable, but I haven't had a chance to go through it in detail (and likely won't for a while)
I'd say anything we can do to limit scope, with the goal of limiting complexity, is worth considering.

TomAugspurger · 2025-08-30T22:16:51Z

src/zarr/abc/codec.py

-class CodecJSON_V2(TypedDict, Generic[TName]):
-    """The JSON representation of a codec for Zarr V2"""


Probably should reexport these types here for backwards compatibility.

done in 30d48a8

TomAugspurger · 2025-08-30T22:20:47Z

src/zarr/core/array.py

            If the dictionary data is invalid or incompatible with either Zarr format 2 or 3 array creation.
        """
-        metadata = parse_array_metadata(data)
+        metadata = parse_array_metadata(data)  # type: ignore[call-overload]


Is the type: ignore needed due to the presence of ArrayMetadata in the signature for __init__?

we need the overload because the data input to from_dict is typed as dict[str, JSON]. I could change that to the union of the two typeddict metadata types, but this will incompatibly override the base class implementation of from_dict, and so I will need another # type: ignore there. The only way to clean all of this up is to redo our base Metadata ABC, which I deemed out of scope for this PR

TomAugspurger · 2025-08-30T22:27:00Z

src/zarr/core/common.py

+    """
+
+    kind: Literal["inline"]
+    must_understand: Literal["false"]


Is this correct for v2 consolidated metadata? IIRC, it didn't have must_understand and maybe had a version field? I guess maybe I'm mixing up consolidated metadata as written by zarr-python 2.x, and the v3 spec.

wasn't v2 consolidated metadata exclusively via an external .zmetadata file, with no change to .zattrs?

you're right, the .zmetadata file had version and "metadata" keys:

out = { "zarr_consolidated_format": 1, "metadata": {key: json_loads(store[key]) for key in store if is_zarr_key(key)}, }

But the above typeddict is for modelling the use of our inline consolidated metadata model for Zarr V2 arrays. I would need a totally separate type for the zmetadata contents.

TomAugspurger · 2025-08-30T22:37:04Z

src/zarr/core/type_check.py

+    return tp
+
+
+def check_type(obj: Any, expected_type: Any, path: str = "value") -> TypeCheckResult:


Can we narrow the type on expected_type here? And do we gain anything if we do?

I'd like to minimize complexity. We know we have a relatively constrained set of expected objects we're going to be passing here (node metadata, codec configuration, ...). If there's anything we can leverage to make this simpler, let's do it.

i think you're kind of asking if we could define a "type-checkable type", i.e. a union of all the types we support runtime type checking for. I'd like that, but I haven't looked into setting that up yet. I don't know what the blockers would be, and I think it would be potentially helpful to more clearly define the scope of this checker.

although, in its current implementation, any type that isn't associated with a specific checking routine falls back to isinstance, which does work for pretty much all user-defined classes...

have a look at c7096b1, It wasn't too painful to narrow from Any to type | types.UnionType | ForwardRef | None

TomAugspurger · 2025-08-30T22:39:18Z

src/zarr/core/type_check.py

+T = TypeVar("T")
+
+
+def ensure_type(obj: object, expected_type: type[T], path: str = "value") -> T:


path is unused?

it's passed to check_type a few lines down

…at/type-checker

d-v-b · 2025-09-01T21:40:09Z

tests/test_metadata/test_consolidated.py

                                        "array": ArrayV3Metadata.from_dict(
                                            {
-                                                **array_metadata,
+                                                **array_metadata,  # type: ignore[typeddict-item]


mypy doesn't know about structural assignability in typeddicts, i.e. the fact that typeddicts tolerate extra keys

d-v-b added 6 commits August 22, 2025 11:25

working type checker

b2f4ff0

working type check that fails for recursive parameters in generic typ…

7adff59

…eddicts

working recursive parameters for generic typeddict

35c7203

all but notrequired

07e2315

switch up imports

85b48df

remove cache poisoing bug, and deploy type checker throughout the cod…

4fe9ae4

…ebase

github-actions bot added the needs release notes Automatically applied to PRs which haven't added release notes label Aug 24, 2025

d-v-b changed the title ~~add a runttime type checker for metadata objects~~ add a runtime type checker for metadata objects Aug 24, 2025

d-v-b added 12 commits August 24, 2025 21:00

restore dimension name normalization

2125153

fix array metadata dicts and refactor to_dict test

32cd309

lint

21d6188

fan out v3 metadata test

6125c1b

update from_dict

cf0615b

overloads for parse_array_metadata

7fce136

Merge branch 'main' of github.com:zarr-developers/zarr-python into fe…

b467747

…at/type-checker

fix missing imports

943e148

add more type information

d1be08c

fix bugs, refine structured data type json representation

ea3ed12

remove unnnecessary test case

a098cc2

changelog

fc06ab4

github-actions bot removed the needs release notes Automatically applied to PRs which haven't added release notes label Aug 25, 2025

d-v-b marked this pull request as ready for review August 25, 2025 13:09

bump minimal typing_extensions version to the release that included e…

1d4bd72

…valuate_forward_ref

d-v-b added 4 commits August 25, 2025 17:00

Merge branch 'main' of github.com:zarr-developers/zarr-python into fe…

d061fe1

…at/type-checker

Merge branch 'main' into feat/type-checker

bbd8ba7

improve error messages, compactify tests, add special case for disamb…

11f7499

…iguating ints from bools

Merge branch 'main' into feat/type-checker

1892df1

d-v-b added 3 commits August 29, 2025 20:40

Merge branch 'main' into feat/type-checker

eda19ec

Merge branch 'main' of github.com:zarr-developers/zarr-python into fe…

bb7e84e

…at/type-checker

remove dead code and consolidate tests

4cc0385

TomAugspurger reviewed Aug 30, 2025

View reviewed changes

d-v-b added 4 commits September 1, 2025 23:18

Merge branch 'main' of github.com:zarr-developers/zarr-python into fe…

be71a87

…at/type-checker

remove redundant imports

971945b

re-export codecjson type

30d48a8

more re-exports

a483c73

d-v-b commented Sep 1, 2025

View reviewed changes

d-v-b added 2 commits September 1, 2025 23:59

narrow input type of type_check

c7096b1

Merge branch 'main' into feat/type-checker

9eb287b

		class CodecJSON_V2(TypedDict, Generic[TName]):
		"""The JSON representation of a codec for Zarr V2"""

		return tp


		def check_type(obj: Any, expected_type: Any, path: str = "value") -> TypeCheckResult:

		T = TypeVar("T")


		def ensure_type(obj: object, expected_type: type[T], path: str = "value") -> T:

Uh oh!

add a runtime type checker for metadata objects #3400

Are you sure you want to change the base?

add a runtime type checker for metadata objects #3400

Uh oh!

Conversation

d-v-b commented Aug 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Demo

Supported types

cost

maintenance burden

performance

benefit

alternatives

Uh oh!

d-v-b commented Aug 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov bot commented Aug 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

d-v-b commented Aug 29, 2025

Uh oh!

d-v-b commented Aug 29, 2025

Uh oh!

TomAugspurger left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

d-v-b commented Aug 24, 2025 •

edited

Loading

d-v-b commented Aug 25, 2025 •

edited

Loading

codecov bot commented Aug 25, 2025 •

edited

Loading