GroupSpec signature

[d-v-b]

Currently we have

TAttr = TypeVar("TAttr", bound=Mapping[str, Any])
TItem = TypeVar("TItem", bound=Union["GroupSpec", "ArraySpec"])

class GroupSpec(NodeSpecV2, Generic[TAttr, TItem]):
    attributes: TAttr
    members: dict[str, TItem]

This excludes the key space of members from the type domain (the keys can be any strings).

What if we did this instead:

TAttr = TypeVar("TAttr", bound=Mapping[str, Any])
TItem = TypeVar("TItem", bound=Dict[str, Union["GroupSpec", "ArraySpec"]])

class GroupSpec(NodeSpecV2, Generic[TAttr, TItem]):
    attributes: TAttr
    members: TItem

I think this would enable requiring that Zarr groups contain nodes with specific names, e.g.

class MyMembers(TypedDict):
    s0: ArraySpec[ArrayAttrs]
    s1: ArraySpec[ArrayAttrs]

GroupSpec[MyAttributes, MyMembers](...)

But it's not clear to how to make this play nice with pydantic models, since they can't be coerced to Dict[str, blablabla]

[d-v-b]

I'm working with earth observation data where the names of the members of a group are an important part of the signature, so I'm very interested in this issue. With the new "closed" parameter for typeddicts, this is now possible:

# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "pydantic-zarr",
#    "pydantic>=2.12",
# ]
# ///
from typing import Any, Generic, Mapping, TypeAlias, TypeVar, Union
from pydantic import BaseModel
from typing_extensions import TypedDict
from pydantic_zarr.v2 import ArraySpec, NodeSpec

TBaseAttr: TypeAlias = Mapping[str, object] | BaseModel
TBaseMember: TypeAlias = Mapping[str, Union["GroupSpec", "ArraySpec"]]

AnyArraySpec: TypeAlias = "ArraySpec[Any]"
AnyGroupSpec: TypeAlias = "GroupSpec[Any, Any]"

TAttr = TypeVar("TAttr", bound=TBaseAttr)
TMembers = TypeVar("TMembers", bound=TBaseMember)

class GroupSpec(NodeSpec, Generic[TAttr, TMembers]):
    attributes: TAttr
    members: TMembers

class Foo(TypedDict, closed=True):
    a: GroupSpec[Any, Any]
    b: ArraySpec[Any]

class G(GroupSpec[Any, Foo]):
    ...

aspec = ArraySpec(
    shape=(10,), 
    dtype='uint8', 
    chunks=(1,), 
    dimension_separator='.',
    attributes={})

gspec = GroupSpec(members={}, attributes={})

try:
    wrong = G(members={"a": gspec, "b": aspec, "c": gspec}, attributes={}) # extra key in members
except ValueError as e:
    print(e)
"""
1 validation error for G
members.c
  Extra inputs are not permitted [type=extra_forbidden, input_value=GroupSpec(zarr_format=2, ...tributes={}, members={}), input_type=GroupSpec]
    For further information visit https://errors.pydantic.dev/2.12/v/extra_forbidden
"""
try:
    wrong = G(members={"b": aspec}, attributes={}) # missing key from members
except ValueError as e:
    print(e)
"""
1 validation error for G
members.a
  Field required [type=missing, input_value={'b': ArraySpec(zarr_form...r='.', compressor=None)}, input_type=dict]
    For further information visit https://errors.pydantic.dev/2.12/v/missing
"""
right = G(members={"a": gspec, "b": aspec}, attributes={})
right_2 = G(members={"a": gspec.model_dump(), "b": aspec.model_dump()}, attributes={}) # type checkers hate this!

[d-v-b]

cc @dstansby , I think this change might be good to roll out alongside removing the attributes generic type parameter from ArraySpec / GroupSpec (#96 (comment)), thoughts?

[dstansby]

This sounds good to me, and potentially not breaking? #96 (comment) will be a major breaking change anyway, so I guess it doesn't matter how breaking this is.

[d-v-b]

the change i'm proposing here would be somewhat breaking, because the type of the generic parameter for the members of a group would change. Currently, you do GroupSpec[T, M] where M denotes the type of the values in the members field. I'm proposing GroupSpec[T, M], where M denotes the type of the entire members field, i.e. keys AND values. But the migration would be really simple.

[d-v-b]

to elaborate on how this came up for me: i'm working with some zarr data where groups have no attributes at all, they just have subgroups with different names.

The set of the names of the subgroups effectively defines the type of the group, but with our current type annotations, there's no way to give pydantic access to a model of what the names of the subgroups should be, and so pydantic fails to resolve the zarr groups into the correct pydantic models.