feature(zarr-metadata): partial metadata

[d-v-b]

The array / group metadata classes in zarr-metadata model the required / not required fields in array / group metadata documents. But there are situations where it's useful to mark all fields as not required, e.g. for a function update_metadata(old_metadata: ArrayMetadataV3, changes: ArrayMetadataUpdateV3) -> ArrayMetadataV3, or for cases (often in tests) where you only need to construct a subset of a metadata document.

So zarr-metadata should add variants of the core metadata documents where all fields are NotRequired.

[chuckwondo]

Offhand, rather than define an additional type, I would follow an approach similar to that in zarrs. For example, see https://docs.rs/zarrs/latest/zarrs/array/struct.ArrayMetadataV3.html.

There, the constructor expects values only for all required values (other than the fixed values, such as zarr_format), defaulting all optional values to sensible defaults. An optional value, where the default is not desired, is specified via a corresponding with_* method, resulting in a new object.

Given that here we define ArrayMetadataV3 as a TypedDict, this is not possible. If we were to use a dataclass instead, we could (I think) mimic what zarrs does here.

However, if I understand what you want to achieve with this issue, this can already be satisfied without defining any variants of the metadata types. A TypedDict already permits "updates" to an instance, where we can specify only the entries we wish to change, regardless of which keys are required or not:

old_metadata: ArrayMetadataV3 = ...
new_metadata = ArrayMetadataV3({**old_metadata, "key1": value1, ...})

# This also works (because `dict` allows it), but does not adhere to PEP 589, so fails type checking.
new_metadata = ArrayMetadataV3(old_metadata, key1=value1, ...)

This creates a new object copied from the old object, overriding or adding only the desired key-value pairs, including extra items. (Unfortunately, mypy still does not support extra_items and closed, so we must still ignore its complaints.)

[d-v-b]

what a partial type enables, and the full type with required fields does not support, is exactly the signature I referenced:

def update_metadata(old_metadata: ArrayMetadataV3, changes: ArrayMetadataUpdateV3) -> ArrayMetadataV3

without a partial type, we can't model a type-safe set of changes applied to an existing metadata document. zarrs handles updates with a variety of with_<field> methods, which take a value with the same type as the corresponding field. we don't need a per-field method if we have the type i'm proposing here. we would just have a single method, update(self, changes: ArrayMetadataUpdateV3) -> Self

[d-v-b]

new_metadata = ArrayMetadataV3({**old_metadata, "key1": value1, ...})

this does look appealing if we are just sticking with complete typeddicts. but there are places in our tests where we work with partial metadata documents, and it would be good to make this pattern type safe. ideally python would have something like typescript's PIck utility type...

[chuckwondo]

I'm curious to understand the cases where you believe it would make sense to support partial metadata documents. Can you point me to some specific examples in the tests?

As far as I can tell, this is not possible in zarrs, but please correct me if I'm wrong. If I'm correct, what makes you believe it should be possible here?

I realize that Python's type system is not as sophisticated as Rust's and TypeScript's are, so we are often limited in what we can do with it, so I'm just playing devil's advocate a bit to both help me get a better understanding of this library as well as push towards the most ergonomic typing we can get within Python's limitations.

The more I can understand the various use cases you have in mind, the better I can make sensible suggestions.

[d-v-b]

I'm curious to understand the cases where you believe it would make sense to support partial metadata documents. Can you point me to some specific examples in the tests?

see

zarr-python/tests/test_metadata/test_v3.py

Line 260 in 84532ca

def test_array_metadata_roundtrip(case: Expect[dict[str, Any], dict[str, Any]]) -> None:

for a representative example. In this test, we parametrize over a subset of fields in a metadata document, with the goal of ensuring that our dataclass model of metadata serializes those fields faithfully. The input to the test -- a subset of the key: value pairs that constitute a metadata document -- should be annotated as a partial metadata document, but we don't have that type so we settle for dict[str, Any].

basically any place where this test utility function is used:

zarr-python/tests/test_metadata/conftest.py

Line 11 in 84532ca

def minimal_metadata_dict_v3(

is a place where we would benefit from a partial metadata document type. That function itself is missing such a type from its signature.

it would probably also help if I sketched out where I'd like to go with zarr-python's metadata classes. the metadata classes in zarr-python today ArrayV3Metadata and ArrayV2Metadata do 2 things: they handle JSON ser/deserialization, and they bundle classes used for runtime stuff (codecs, data types, chunk key encodings). As we learned when adding the rectilinear chunks, the JSON metadata and the runtime classes are two levels of detail that probably should not mix.

So I would like to make ArrayV3Metadata simpler: it should still be a plain dataclass (because this is more ergonomic than a typeddict), but the type of the fields should be simple JSON-like types. This makes JSON serialization very simple. Such a dataclass could have a .update method, and the signature of that method would require the partial metadata type to be useful. Then it's the Array class' job to interpret a JSON-flavored metadata document to create whatever runtime classes it needs.

An advantage of a single update method vs multiple with_<field> methods is that there are logically coupled fields: you can't update the shape of an array metadata document without also changing the chunk grid, and possibly the codecs. Updating the data type might also require updating the codecs. A single method can handle all of these cases more easily than a separate method per field.