Repository: astropenguin/xarray-dataclasses

Issue:
Add v2 typing module
Add a new typing module for v2 based on that of pandas-dataclasses (v0.9.0).
Note that it will be included in the current package, but not used until the v2 release.

Comments:

Patch:
diff --git a/xarray_dataclasses/v2/__init__.py b/xarray_dataclasses/v2/__init__.py
new file mode 100644
index 0000000..1f6e4f1
--- /dev/null
+++ b/xarray_dataclasses/v2/__init__.py
@@ -0,0 +1,4 @@
+__all__ = ["typing"]
+
+
+from . import typing
diff --git a/xarray_dataclasses/v2/typing.py b/xarray_dataclasses/v2/typing.py
new file mode 100644
index 0000000..e4865a3
--- /dev/null
+++ b/xarray_dataclasses/v2/typing.py
@@ -0,0 +1,216 @@
+__all__ = ["Attr", "Coord", "Coordof", "Data", "Dataof", "Other"]
+
+
+# standard library
+from dataclasses import Field
+from enum import Enum, auto
+from itertools import chain
+from typing import (
+    Any,
+    Callable,
+    Collection,
+    Dict,
+    Generic,
+    Hashable,
+    Iterable,
+    Optional,
+    Tuple,
+    TypeVar,
+    Union,
+)
+
+
+# dependencies
+import numpy as np
+import xarray as xr
+from typing_extensions import (
+    Annotated,
+    Literal,
+    ParamSpec,
+    Protocol,
+    get_args,
+    get_origin,
+    get_type_hints,
+)
+
+
+# type hints (private)
+P = ParamSpec("P")
+T = TypeVar("T")
+TDataClass = TypeVar("TDataClass", bound="DataClass[Any]")
+TDataArray = TypeVar("TDataArray", bound=xr.DataArray)
+TDataset = TypeVar("TDataset", bound=xr.Dataset)
+TDims = TypeVar("TDims")
+TDType = TypeVar("TDType")
+TXarray = TypeVar("TXarray", bound="Xarray")
+Xarray = Union[xr.DataArray, xr.Dataset]
+
+
+class DataClass(Protocol[P]):
+    """Type hint for dataclass objects."""
+
+    __dataclass_fields__: Dict[str, "Field[Any]"]
+
+    def __init__(self, *args: P.args, **kwargs: P.kwargs) -> None:
+        ...
+
+
+class XarrayClass(Protocol[P, TXarray]):
+    """Type hint for dataclass objects with a xarray factory."""
+
+    __dataclass_fields__: Dict[str, "Field[Any]"]
+    __xarray_factory__: Callable[..., TXarray]
+
+    def __init__(self, *args: P.args, **kwargs: P.kwargs) -> None:
+        ...
+
+
+class Dims(Generic[TDims]):
+    """Empty class for storing type of dimensions."""
+
+    pass
+
+
+class Role(Enum):
+    """Annotations for typing dataclass fields."""
+
+    ATTR = auto()
+    """Annotation for attribute fields."""
+
+    COORD = auto()
+    """Annotation for coordinate fields."""
+
+    DATA = auto()
+    """Annotation for data fields."""
+
+    OTHER = auto()
+    """Annotation for other fields."""
+
+    @classmethod
+    def annotates(cls, tp: Any) -> bool:
+        """Check if any role annotates a type hint."""
+        return any(isinstance(arg, cls) for arg in get_args(tp))
+
+
+# type hints (public)
+Attr = Annotated[T, Role.ATTR]
+"""Type hint for attribute fields (``Attr[T]``)."""
+
+Coord = Annotated[Union[Dims[TDims], Collection[TDType]], Role.COORD]
+"""Type hint for coordinate fields (``Coord[TDims, TDType]``)."""
+
+Coordof = Annotated[TDataClass, Role.COORD]
+"""Type hint for coordinate fields (``Dataof[TDataClass]``)."""
+
+Data = Annotated[Union[Dims[TDims], Collection[TDType]], Role.DATA]
+"""Type hint for data fields (``Coord[TDims, TDType]``)."""
+
+Dataof = Annotated[TDataClass, Role.DATA]
+"""Type hint for data fields (``Dataof[TDataClass]``)."""
+
+Other = Annotated[T, Role.OTHER]
+"""Type hint for other fields (``Other[T]``)."""
+
+
+# runtime functions
+def deannotate(tp: Any) -> Any:
+    """Recursively remove annotations in a type hint."""
+
+    class Temporary:
+        __annotations__ = dict(tp=tp)
+
+    return get_type_hints(Temporary)["tp"]
+
+
+def find_annotated(tp: Any) -> Iterable[Any]:
+    """Generate all annotated types in a type hint."""
+    args = get_args(tp)
+
+    if get_origin(tp) is Annotated:
+        yield tp
+        yield from find_annotated(args[0])
+    else:
+        yield from chain(*map(find_annotated, args))
+
+
+def get_annotated(tp: Any) -> Any:
+    """Extract the first role-annotated type."""
+    for annotated in filter(Role.annotates, find_annotated(tp)):
+        return deannotate(annotated)
+
+    raise TypeError("Could not find any role-annotated type.")
+
+
+def get_annotations(tp: Any) -> Tuple[Any, ...]:
+    """Extract annotations of the first role-annotated type."""
+    for annotated in filter(Role.annotates, find_annotated(tp)):
+        return get_args(annotated)[1:]
+
+    raise TypeError("Could not find any role-annotated type.")
+
+
+def get_dims(tp: Any) -> Optional[Tuple[str, ...]]:
+    """Extract dimensions if found or return None."""
+    try:
+        dims = get_args(get_args(get_annotated(tp))[0])[0]
+    except (IndexError, TypeError):
+        return None
+
+    args = get_args(dims)
+    origin = get_origin(dims)
+
+    if args == () or args == ((),):
+        return ()
+
+    if origin is Literal:
+        return (str(args[0]),)
+
+    if not (origin is tuple or origin is Tuple):
+        raise TypeError(f"Could not find any dims in {tp!r}.")
+
+    if not all(get_origin(arg) is Literal for arg in args):
+        raise TypeError(f"Could not find any dims in {tp!r}.")
+
+    return tuple(str(get_args(arg)[0]) for arg in args)
+
+
+def get_dtype(tp: Any) -> Optional[str]:
+    """Extract a data type if found or return None."""
+    try:
+        dtype = get_args(get_args(get_annotated(tp))[1])[0]
+    except (IndexError, TypeError):
+        return None
+
+    if dtype is Any or dtype is type(None):
+        return None
+
+    if get_origin(dtype) is Literal:
+        dtype = get_args(dtype)[0]
+
+    return np.dtype(dtype).name
+
+
+def get_name(tp: Any, default: Hashable = None) -> Hashable:
+    """Extract a name if found or return given default."""
+    try:
+        name = get_annotations(tp)[1]
+    except (IndexError, TypeError):
+        return default
+
+    if name is Ellipsis:
+        return default
+
+    try:
+        hash(name)
+    except TypeError:
+        raise ValueError("Could not find any valid name.")
+
+    return name
+
+
+def get_role(tp: Any, default: Role = Role.OTHER) -> Role:
+    """Extract a role if found or return given default."""
+    try:
+        return get_annotations(tp)[0]
+    except TypeError:
+        return default


Test patch:
diff --git a/tests/test_v2_typing.py b/tests/test_v2_typing.py
new file mode 100644
index 0000000..e8b177f
--- /dev/null
+++ b/tests/test_v2_typing.py
@@ -0,0 +1,139 @@
+# standard library
+from dataclasses import dataclass
+from typing import Any, Tuple, Union
+
+
+# dependencies
+import numpy as np
+from xarray_dataclasses.v2.typing import (
+    Attr,
+    Coord,
+    Coordof,
+    Data,
+    Dataof,
+    Role,
+    get_dims,
+    get_dtype,
+    get_name,
+    get_role,
+)
+from pytest import mark
+from typing_extensions import Annotated as Ann, Literal as L
+
+
+# test data
+@dataclass
+class DataClass:
+    data: Any
+
+
+testdata_dims = [
+    (Coord[Tuple[()], Any], ()),
+    (Coord[L["x"], Any], ("x",)),
+    (Coord[Tuple[L["x"]], Any], ("x",)),
+    (Coord[Tuple[L["x"], L["y"]], Any], ("x", "y")),
+    (Coordof[DataClass], None),
+    (Data[Tuple[()], Any], ()),
+    (Data[L["x"], Any], ("x",)),
+    (Data[Tuple[L["x"]], Any], ("x",)),
+    (Data[Tuple[L["x"], L["y"]], Any], ("x", "y")),
+    (Dataof[DataClass], None),
+    (Ann[Coord[L["x"], Any], "coord"], ("x",)),
+    (Ann[Coordof[DataClass], "coord"], None),
+    (Ann[Data[L["x"], Any], "data"], ("x",)),
+    (Ann[Dataof[DataClass], "data"], None),
+    (Union[Ann[Coord[L["x"], Any], "coord"], Ann[Any, "any"]], ("x",)),
+    (Union[Ann[Coordof[DataClass], "coord"], Ann[Any, "any"]], None),
+    (Union[Ann[Data[L["x"], Any], "data"], Ann[Any, "any"]], ("x",)),
+    (Union[Ann[Dataof[DataClass], "data"], Ann[Any, "any"]], None),
+]
+
+testdata_dtype = [
+    (Coord[Any, Any], None),
+    (Coord[Any, None], None),
+    (Coord[Any, int], np.dtype("i8")),
+    (Coord[Any, L["i8"]], np.dtype("i8")),
+    (Coordof[DataClass], None),
+    (Data[Any, Any], None),
+    (Data[Any, None], None),
+    (Data[Any, int], np.dtype("i8")),
+    (Data[Any, L["i8"]], np.dtype("i8")),
+    (Dataof[DataClass], None),
+    (Ann[Coord[Any, float], "coord"], np.dtype("f8")),
+    (Ann[Coordof[DataClass], "coord"], None),
+    (Ann[Data[Any, float], "data"], np.dtype("f8")),
+    (Ann[Dataof[DataClass], "data"], None),
+    (Union[Ann[Coord[Any, float], "coord"], Ann[Any, "any"]], np.dtype("f8")),
+    (Union[Ann[Coordof[DataClass], "coord"], Ann[Any, "any"]], None),
+    (Union[Ann[Data[Any, float], "data"], Ann[Any, "any"]], np.dtype("f8")),
+    (Union[Ann[Dataof[DataClass], "data"], Ann[Any, "any"]], None),
+]
+
+testdata_name = [
+    (Attr[Any], None),
+    (Coord[Any, Any], None),
+    (Coordof[DataClass], None),
+    (Data[Any, Any], None),
+    (Dataof[DataClass], None),
+    (Any, None),
+    (Ann[Attr[Any], "attr"], "attr"),
+    (Ann[Coord[Any, Any], "coord"], "coord"),
+    (Ann[Coordof[DataClass], "coord"], "coord"),
+    (Ann[Data[Any, Any], "data"], "data"),
+    (Ann[Dataof[DataClass], "data"], "data"),
+    (Ann[Any, "other"], None),
+    (Ann[Attr[Any], ..., "attr"], None),
+    (Ann[Coord[Any, Any], ..., "coord"], None),
+    (Ann[Coordof[DataClass], ..., "coord"], None),
+    (Ann[Data[Any, Any], ..., "data"], None),
+    (Ann[Dataof[DataClass], ..., "data"], None),
+    (Ann[Any, ..., "other"], None),
+    (Union[Ann[Attr[Any], "attr"], Ann[Any, "any"]], "attr"),
+    (Union[Ann[Coord[Any, Any], "coord"], Ann[Any, "any"]], "coord"),
+    (Union[Ann[Coordof[DataClass], "coord"], Ann[Any, "any"]], "coord"),
+    (Union[Ann[Data[Any, Any], "data"], Ann[Any, "any"]], "data"),
+    (Union[Ann[Dataof[DataClass], "data"], Ann[Any, "any"]], "data"),
+    (Union[Ann[Any, "other"], Ann[Any, "any"]], None),
+]
+
+testdata_role = [
+    (Attr[Any], Role.ATTR),
+    (Coord[Any, Any], Role.COORD),
+    (Coordof[DataClass], Role.COORD),
+    (Data[Any, Any], Role.DATA),
+    (Dataof[DataClass], Role.DATA),
+    (Any, Role.OTHER),
+    (Ann[Attr[Any], "attr"], Role.ATTR),
+    (Ann[Coord[Any, Any], "coord"], Role.COORD),
+    (Ann[Coordof[DataClass], "coord"], Role.COORD),
+    (Ann[Data[Any, Any], "data"], Role.DATA),
+    (Ann[Dataof[DataClass], "data"], Role.DATA),
+    (Ann[Any, "other"], Role.OTHER),
+    (Union[Ann[Attr[Any], "attr"], Ann[Any, "any"]], Role.ATTR),
+    (Union[Ann[Coord[Any, Any], "coord"], Ann[Any, "any"]], Role.COORD),
+    (Union[Ann[Coordof[DataClass], "coord"], Ann[Any, "any"]], Role.COORD),
+    (Union[Ann[Data[Any, Any], "data"], Ann[Any, "any"]], Role.DATA),
+    (Union[Ann[Dataof[DataClass], "data"], Ann[Any, "any"]], Role.DATA),
+    (Union[Ann[Any, "other"], Ann[Any, "any"]], Role.OTHER),
+]
+
+
+# test functions
+@mark.parametrize("tp, dims", testdata_dims)
+def test_get_dims(tp: Any, dims: Any) -> None:
+    assert get_dims(tp) == dims
+
+
+@mark.parametrize("tp, dtype", testdata_dtype)
+def test_get_dtype(tp: Any, dtype: Any) -> None:
+    assert get_dtype(tp) == dtype
+
+
+@mark.parametrize("tp, name", testdata_name)
+def test_get_name(tp: Any, name: Any) -> None:
+    assert get_name(tp) == name
+
+
+@mark.parametrize("tp, role", testdata_role)
+def test_get_role(tp: Any, role: Any) -> None:
+    assert get_role(tp) is role