Repository: astropenguin/xarray-dataclasses (python)

Issue:
Fix parsing dataclass under __future__.annotations enabled
The current (v0.4.0) dataclass parser (`xarray_dataclasses.parse`) does not work when the postponed evaluation of annotations ([PEP 563](https://www.python.org/dev/peps/pep-0563)) is enabled by importing `from __future__ import annotations`. This is because the parser gets type hints of dataclass fields from `cls.__dataclass_fields__`: Under such environments, `field.type` becomes string and loses runtime information.

This issue is for resolving the problem by using `typing_extensions.get_type_hints(cls, include_extras=True)` that ensures evaluated type hints made from strings.

Pull request:
Fix parsing dataclass under __future__.annotations enabled
This PR closes #65.

Patch:
diff --git a/xarray_dataclasses/datamodel.py b/xarray_dataclasses/datamodel.py
index 70e08e7..ac65410 100644
--- a/xarray_dataclasses/datamodel.py
+++ b/xarray_dataclasses/datamodel.py
@@ -9,16 +9,18 @@ from typing import Any, Callable, cast, Generic, List, TypeVar, Union
 # dependencies
 import numpy as np
 import xarray as xr
-from typing_extensions import get_args, TypedDict
+from typing_extensions import TypedDict
 
 
 # submodules
+from .deprecated import get_type_hints
 from .typing import (
     ArrayLike,
     DataClass,
     Dims,
     Dtype,
     FieldType,
+    get_class,
     get_dims,
     get_dtype,
     Reference,
@@ -80,10 +82,8 @@ class Data(FieldModel[xr.DataArray]):
     @classmethod
     def from_field(cls, field: Field[Any], value: Any) -> "Data":
         """Create a field model from a dataclass field and a value."""
-        args = get_args(get_args(unannotate(field.type))[0])
-
-        dims = get_dims(args[0])
-        dtype = get_dtype(args[1])
+        dims = get_dims(field.type)
+        dtype = get_dtype(field.type)
         type: DataArrayDict = {"dims": dims, "dtype": dtype}
 
         def factory(value: Any, reference: Reference) -> xr.DataArray:
@@ -102,7 +102,7 @@ class Dataof(FieldModel[xr.DataArray]):
     @classmethod
     def from_field(cls, field: Field[Any], value: Any) -> "Dataof":
         """Create a field model from a dataclass field and a value."""
-        dataclass = get_args(unannotate(field.type))[0]
+        dataclass = get_class(field.type)
         type = resolve_class(dataclass)
 
         def factory(value: Any, reference: Reference) -> xr.DataArray:
@@ -151,6 +151,7 @@ class DataModel:
     def from_dataclass(cls, dataclass: DataClass) -> "DataModel":
         """Create a data model from a dataclass or its object."""
         model = cls()
+        eval_field_types(dataclass)
 
         for field_ in dataclass.__dataclass_fields__.values():
             value = getattr(dataclass, field_.name, field_.default)
@@ -172,6 +173,15 @@ class DataModel:
 
 
 # runtime functions
+def eval_field_types(dataclass: DataClass) -> None:
+    """Evaluate field types of a dataclass or its object."""
+    hints = get_type_hints(dataclass, include_extras=True)  # type: ignore
+
+    for field_ in dataclass.__dataclass_fields__.values():
+        if isinstance(field_.type, str):
+            field_.type = hints[field_.name]
+
+
 def typedarray(
     data: Any,
     dims: Dims,
diff --git a/xarray_dataclasses/deprecated.py b/xarray_dataclasses/deprecated.py
index 59d0248..aef8069 100644
--- a/xarray_dataclasses/deprecated.py
+++ b/xarray_dataclasses/deprecated.py
@@ -1,20 +1,22 @@
 """Submodule for functions that will be deprecated in v1.0.0."""
-__all__ = ["dataarrayclass", "datasetclass"]
+__all__ = ["dataarrayclass", "datasetclass", "get_type_hints"]
 
 
 # standard library
+import re
 from dataclasses import dataclass, Field
-from typing import Any, Callable, Dict, Optional, Type, TypeVar, Union
+from typing import Any, Callable, Dict, ForwardRef, Optional, Type, TypeVar, Union
+from typing import _eval_type  # type: ignore
 from warnings import warn
 
 
 # dependencies
-from typing_extensions import Protocol
+from typing_extensions import Literal, Protocol
+from typing_extensions import get_type_hints as _get_type_hints
 
 
-# submodules
-from .dataarray import AsDataArray
-from .dataset import AsDataset
+# constants
+GENERIC_NAME = re.compile(r"\[.+\]")
 
 
 # type hints
@@ -41,13 +43,13 @@ def dataarrayclass(
     shorthands: bool = True,
 ) -> Union[Type[DataClass], Callable[[type], Type[DataClass]]]:
     """Class decorator to create a DataArray class."""
+    from .dataarray import AsDataArray
 
     warn(
-        DeprecationWarning(
-            "This decorator will be removed in v1.0.0. ",
-            "Please consider to use the Python's dataclass ",
-            "and the mix-in class (AsDataArray) instead.",
-        )
+        "This decorator will be removed in v1.0.0. "
+        "Please consider to use the Python's dataclass "
+        "and the mix-in class (AsDataArray) instead.",
+        category=FutureWarning,
     )
 
     def to_dataclass(cls: Type[Any]) -> Type[DataClass]:
@@ -81,13 +83,13 @@ def datasetclass(
     shorthands: bool = True,
 ) -> Union[Type[DataClass], Callable[[type], Type[DataClass]]]:
     """Class decorator to create a Dataset class."""
+    from .dataset import AsDataset
 
     warn(
-        DeprecationWarning(
-            "This decorator will be removed in v1.0.0. ",
-            "Please consider to use the Python's dataclass ",
-            "and the mix-in class (AsDataset) instead.",
-        )
+        "This decorator will be removed in v1.0.0. "
+        "Please consider to use the Python's dataclass "
+        "and the mix-in class (AsDataset) instead.",
+        category=FutureWarning,
     )
 
     def to_dataclass(cls: Type[Any]) -> Type[DataClass]:
@@ -109,6 +111,27 @@ def datasetclass(
         return to_dataclass(cls)
 
 
+def eval_type(type: Any, *args: Any, **kwargs: Any) -> Any:
+    if not isinstance(type, ForwardRef):
+        return _eval_type(type, *args, **kwargs)
+
+    name = type.__forward_arg__
+
+    if GENERIC_NAME.search(name):
+        return _eval_type(type, *args, **kwargs)
+
+    warn(
+        f"For backward compatibility, forward reference {name!r} "
+        f"was not evaluated but converted to Literal[{name!r}]. "
+        "From v1.0.0, it will be evaluated by default, which may "
+        f"raise NameError or unexpectedly assign an object to {name}. "
+        f"Please consider to replace {name!r} with Literal[{name!r}].",
+        category=FutureWarning,
+    )
+
+    return Literal[name]
+
+
 def extend_class(cls: Type[T], mixin: Type[Any]) -> Type[T]:
     """Extend a class with a mix-in class."""
     if cls.__bases__ == (object,):
@@ -117,3 +140,25 @@ def extend_class(cls: Type[T], mixin: Type[Any]) -> Type[T]:
         bases = (*cls.__bases__, mixin)
 
     return type(cls.__name__, bases, cls.__dict__.copy())
+
+
+def get_type_hints(
+    obj: Callable[..., Any],
+    globalns: Optional[Dict[str, Any]] = None,
+    localns: Optional[Dict[str, Any]] = None,
+    include_extras: bool = False,
+) -> Dict[str, Any]:
+    """Return type hints for an object.
+
+    Unlike the original (``typing.get_type_hints``), it does NOT handle
+    forward references encoded as string but converts them to literal types.
+    Other behavior is the same as the original.
+
+    """
+    import typing
+
+    try:
+        typing._eval_type = eval_type  # type: ignore
+        return _get_type_hints(obj, globalns, localns, include_extras)
+    finally:
+        typing._eval_type = _eval_type  # type: ignore
diff --git a/xarray_dataclasses/typing.py b/xarray_dataclasses/typing.py
index 9a40890..ad80865 100644
--- a/xarray_dataclasses/typing.py
+++ b/xarray_dataclasses/typing.py
@@ -4,16 +4,7 @@ __all__ = ["Attr", "Coord", "Coordof", "Data", "Dataof", "Name"]
 # standard library
 from dataclasses import Field
 from enum import auto, Enum
-from itertools import chain
-from typing import (
-    Any,
-    Dict,
-    ForwardRef,
-    Optional,
-    Tuple,
-    TypeVar,
-    Union,
-)
+from typing import Any, Dict, Optional, Tuple, TypeVar, Union
 
 
 # dependencies
@@ -22,6 +13,7 @@ from typing_extensions import (
     Annotated,
     get_args,
     get_origin,
+    get_type_hints,
     Literal,
     Protocol,
     runtime_checkable,
@@ -277,80 +269,71 @@ Examples:
 
 
 # runtime functions
-def get_dims(obj: Any) -> Dims:
-    """Parse an object to get dims."""
-    obj = unannotate(obj)
-
-    if obj == ():
-        return ()
+def get_class(hint: Any) -> Any:
+    """Return a class parsed from a type hint."""
+    return get_first(unannotate(hint))
 
-    if obj is NoneType:
-        return ()
 
-    if isinstance(obj, str):
-        return (obj,)
+def get_dims(hint: Any) -> Dims:
+    """Return dims parsed from a type hint."""
+    t_dims = get_args(get_class(hint))[0]
 
-    if isinstance(obj, ForwardRef):
-        return (obj.__forward_arg__,)
+    if is_str_literal(t_dims):
+        return (get_first(t_dims),)
 
-    args = get_args(obj)
-    origin = get_origin(obj)
+    args = get_args(t_dims)
 
-    if origin is tuple:
-        return tuple(chain(*map(get_dims, args)))
+    if args == () or args == ((),):
+        return ()
 
-    if origin is Literal and len(args) == 1:
-        return (str(args[0]),)
+    if all(map(is_str_literal, args)):
+        return tuple(map(get_first, args))
 
-    raise ValueError(f"Could not parse {obj} as dims.")
+    raise ValueError(f"Could not parse dims from {hint!r}.")
 
 
-def get_dtype(obj: Any) -> Dtype:
-    """Parse an object to get dtype."""
-    obj = unannotate(obj)
+def get_dtype(hint: Any) -> Dtype:
+    """Return dtype parsed from a type hint."""
+    t_dtype = get_args(get_class(hint))[1]
 
-    if obj is Any:
+    if t_dtype is Any:
         return None
 
-    if obj is NoneType:
+    if t_dtype is NoneType:
         return None
 
-    if isinstance(obj, str):
-        return obj
+    if isinstance(t_dtype, type):
+        return t_dtype.__name__
+
+    if is_str_literal(t_dtype):
+        return get_first(t_dtype)
 
-    if isinstance(obj, type):
-        return obj.__name__
+    raise ValueError(f"Could not parse dtype from {hint!r}.")
 
-    if isinstance(obj, ForwardRef):
-        return obj.__forward_arg__
 
-    args = get_args(obj)
-    origin = get_origin(obj)
+def get_first(hint: Any) -> Any:
+    """Return the first argument in a type hint."""
+    return get_args(hint)[0]
 
-    if origin is Literal and len(args) == 1:
-        return str(args[0])
 
-    raise ValueError(f"Could not parse {obj} as dtype.")
+def is_str_literal(hint: Any) -> bool:
+    """Check if a type hint is Literal[str]."""
+    args = get_args(hint)
+    origin = get_origin(hint)
 
+    if origin is not Literal:
+        return False
 
-def unannotate(obj: T) -> T:
-    """Recursively remove Annotated types."""
-    import typing
+    if not len(args) == 1:
+        return False
 
-    args = get_args(obj)
-    origin = get_origin(obj)
+    return isinstance(args[0], str)
 
-    if origin is None:
-        return obj
 
-    if origin is Annotated:
-        return unannotate(args[0])
+def unannotate(hint: Any) -> Any:
+    """Recursively remove Annotated type hints."""
 
-    args = map(unannotate, args)
-    args = tuple(filter(None, args))
+    class Temp:
+        __annotations__ = dict(hint=hint)
 
-    try:
-        return origin[args]
-    except TypeError:
-        name = origin.__name__.capitalize()
-        return getattr(typing, name)[args]
+    return get_type_hints(Temp)["hint"]


Test patch:
diff --git a/tests/test_typing.py b/tests/test_typing.py
index 497a2ce..53f2e88 100644
--- a/tests/test_typing.py
+++ b/tests/test_typing.py
@@ -1,60 +1,44 @@
 # standard library
-from typing import Any, ForwardRef, Tuple, Union
+from typing import Any, Tuple
 
 
 # third-party packages
 from pytest import mark
-from typing_extensions import Annotated, Literal
+from typing_extensions import Literal
 
 
 # submodules
-from xarray_dataclasses.typing import get_dims, get_dtype, unannotate
+from xarray_dataclasses.typing import Data, get_dims, get_dtype
 
 
 # type hints
-x = ForwardRef("x")
-y = ForwardRef("y")
+Int = Literal["int"]
 X = Literal["x"]
 Y = Literal["y"]
 
 
 # test datasets
 testdata_dims = [
-    (Tuple[()], ()),
-    (X, ("x",)),
-    (x, ("x",)),
-    (Tuple[X], ("x",)),
-    (Tuple[x], ("x",)),  # type: ignore
-    (Tuple[X, Y], ("x", "y")),
-    (Tuple[x, y], ("x", "y")),  # type: ignore
+    (Data[X, Any], ("x",)),
+    (Data[Tuple[()], Any], ()),
+    (Data[Tuple[X], Any], ("x",)),
+    (Data[Tuple[X, Y], Any], ("x", "y")),
 ]
 
 testdata_dtype = [
-    (Any, None),
-    (int, "int"),
-    ("int", "int"),
-    (Literal["int"], "int"),
-]
-
-testdata_unannotate = [
-    (x, x),
-    (X, X),
-    (Annotated[X, 0], X),
-    (Union[Annotated[X, 0], Y], Union[X, Y]),
+    (Data[X, Any], None),
+    (Data[X, None], None),
+    (Data[X, int], "int"),
+    (Data[X, Int], "int"),
 ]
 
 
 # test functions
-@mark.parametrize("t_dims, dims", testdata_dims)
-def test_get_dims(t_dims: Any, dims: Any) -> None:
-    assert get_dims(t_dims) == dims
-
-
-@mark.parametrize("t_dtype, dtype", testdata_dtype)
-def test_get_dtype(t_dtype: Any, dtype: Any) -> None:
-    assert get_dtype(t_dtype) == dtype
+@mark.parametrize("hint, dims", testdata_dims)
+def test_get_dims(hint: Any, dims: Any) -> None:
+    assert get_dims(hint) == dims
 
 
-@mark.parametrize("t_before, t_after", testdata_unannotate)
-def test_unannotate(t_before: Any, t_after: Any) -> None:
-    assert unannotate(t_before) == t_after
+@mark.parametrize("hint, dtype", testdata_dtype)
+def test_get_dtype(hint: Any, dtype: Any) -> None:
+    assert get_dtype(hint) == dtype