Repository: astropenguin/pandas-dataclasses

Issue:
Update typing module
- [x] Support union in data and index fields
- [x] Support [pandas extended data types](https://pandas.pydata.org/docs/reference/arrays.html)

Discussion:

Patch:
diff --git a/pandas_dataclasses/typing.py b/pandas_dataclasses/typing.py
index 89f913e..0741001 100644
--- a/pandas_dataclasses/typing.py
+++ b/pandas_dataclasses/typing.py
@@ -4,11 +4,23 @@ __all__ = ["Attr", "Data", "Index", "Name", "Other"]
 # standard library
 from dataclasses import Field
 from enum import Enum
-from typing import Any, ClassVar, Collection, Dict, Hashable, Optional, TypeVar, Union
+from typing import (
+    Any,
+    ClassVar,
+    Collection,
+    Dict,
+    Hashable,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+)
 
 
 # dependencies
-import numpy as np
+from numpy import dtype
+from pandas.api.extensions import ExtensionDtype
+from pandas.api.types import pandas_dtype  # type: ignore
 from typing_extensions import (
     Annotated,
     Literal,
@@ -21,12 +33,19 @@ from typing_extensions import (
 
 
 # type hints (private)
-AnyDType: TypeAlias = "np.dtype[Any]"
+AnyDType: TypeAlias = Union["dtype[Any]", ExtensionDtype]
 AnyField: TypeAlias = "Field[Any]"
 T = TypeVar("T")
+TCovariant = TypeVar("TCovariant", covariant=True)
 THashable = TypeVar("THashable", bound=Hashable)
 
 
+class Collection(Collection[TCovariant], Protocol):
+    """Type hint equivalent to typing.Collection."""
+
+    pass
+
+
 class DataClass(Protocol):
     """Type hint for dataclass objects."""
 
@@ -79,28 +98,34 @@ def deannotate(tp: Any) -> Any:
     return get_type_hints(Temporary)["type"]
 
 
-def get_dtype(tp: Any) -> Optional[AnyDType]:
-    """Extract a dtype (NumPy data type) from a type hint."""
+def get_collection(tp: Any) -> Type[Collection[Any]]:
+    """Extract the first collection type from a type hint."""
     tp = deannotate(tp)
 
     if get_origin(tp) is not Union:
-        raise TypeError(f"{tp!r} is not arrayable.")
+        raise TypeError(f"{tp!r} was not a union type.")
 
-    try:
-        tp_array, tp_scalar = get_args(tp)
-    except ValueError:
-        raise TypeError(f"{tp!r} is not arrayable.")
+    # flatten union type after deannotation
+    tp = Union[get_args(tp)]  # type: ignore
 
-    if get_args(tp_array)[0] is not tp_scalar:
-        raise TypeError(f"{tp!r} is not arrayable.")
+    for arg in get_args(tp):
+        if get_origin(arg) is Collection:
+            return arg
+
+    raise TypeError(f"{tp!r} had no collection type.")
+
+
+def get_dtype(tp: Any) -> Optional[AnyDType]:
+    """Extract a dtype (the first data type) from a type hint."""
+    dtype = get_args(get_collection(tp))[0]
 
-    if tp_scalar is Any or tp_scalar is type(None):
-        return None
+    if dtype is Any or dtype is type(None):
+        return
 
-    if get_origin(tp_scalar) is Literal:
-        tp_scalar = get_args(tp_scalar)[0]
+    if get_origin(dtype) is Literal:
+        dtype = get_args(dtype)[0]
 
-    return np.dtype(tp_scalar)
+    return pandas_dtype(dtype)  # type: ignore
 
 
 def get_ftype(tp: Any, default: FType = FType.OTHER) -> FType:


Test patch:
diff --git a/tests/test_typing.py b/tests/test_typing.py
index 5ad02c0..5b039c4 100644
--- a/tests/test_typing.py
+++ b/tests/test_typing.py
@@ -1,9 +1,10 @@
 # standard library
-from typing import Any
+from typing import Any, Optional, Union
 
 
 # dependencies
 import numpy as np
+import pandas as pd
 from pytest import mark
 from typing_extensions import Annotated, Literal
 from pandas_dataclasses.typing import (
@@ -23,10 +24,16 @@ testdata_dtype = [
     (Data[None], None),
     (Data[int], np.dtype("int64")),
     (Data[Literal["i8"]], np.dtype("int64")),
+    (Data[Literal["boolean"]], pd.BooleanDtype()),
     (Index[Any], None),
     (Index[None], None),
     (Index[int], np.dtype("int64")),
     (Index[Literal["i8"]], np.dtype("int64")),
+    (Index[Literal["boolean"]], pd.BooleanDtype()),
+    (Optional[Data[float]], np.dtype("float64")),
+    (Optional[Index[float]], np.dtype("float64")),
+    (Union[Data[float], str], np.dtype("float64")),
+    (Union[Index[float], str], np.dtype("float64")),
 ]
 
 testdata_ftype = [