Repository: astropenguin/pandas-dataclasses

Issue:
Add support for extension array and dtype

Discussion:

Patch:
diff --git a/README.md b/README.md
index 90e72d2..f4cc3c3 100644
--- a/README.md
+++ b/README.md
@@ -408,9 +408,9 @@ where `obj` is a dataclass object that is expected to have `obj.name`.
 
 Release version | Features
 --- | ---
-v0.4.0 | Support for hierarchical column
-v0.5.0 | Support for dynamic naming
-v1.0.0 | Initial major release (freezing public features until v2.0.0)
+v0.5 | Support for dynamic naming
+v0.6 | support for extension array and dtype
+v1.0 | Initial major release (freezing public features until v2.0)
 
 <!-- References -->
 [dataclass]: https://docs.python.org/3/library/dataclasses.html
diff --git a/pandas_dataclasses/core/parsers.py b/pandas_dataclasses/core/parsers.py
index 4b3ebe8..3fbca48 100644
--- a/pandas_dataclasses/core/parsers.py
+++ b/pandas_dataclasses/core/parsers.py
@@ -12,7 +12,7 @@ import pandas as pd
 
 # submodules
 from .specs import Spec
-from .typing import P, AnyDType, DataClass, PandasClass
+from .typing import P, AnyDtype, DataClass, PandasClass
 
 
 # type hints
@@ -93,15 +93,15 @@ def asseries(obj: Any, *, factory: Any = None) -> Any:
     return series
 
 
-def ensure(data: Any, dtype: Optional[AnyDType]) -> Any:
+def ensure(data: Any, dtype: Optional[AnyDtype]) -> Any:
     """Ensure data to be 1D and have given data type."""
-    if not isinstance(data, (pd.Index, pd.Series)):
-        data = np.atleast_1d(data)
+    if not np.ndim(data):
+        data = [data]
 
-    if dtype is None:
-        return data
+    if isinstance(data, (pd.Index, pd.Series)):
+        return type(data)(data, dtype=dtype, copy=False)
     else:
-        return data.astype(dtype, copy=False)
+        return pd.array(data, dtype=dtype, copy=False)
 
 
 def get_attrs(spec: Spec) -> "dict[Hashable, Any]":
diff --git a/pandas_dataclasses/core/specs.py b/pandas_dataclasses/core/specs.py
index a1c5532..7ebc6aa 100644
--- a/pandas_dataclasses/core/specs.py
+++ b/pandas_dataclasses/core/specs.py
@@ -15,7 +15,7 @@ from typing_extensions import Literal, get_type_hints
 # submodules
 from .typing import (
     P,
-    AnyDType,
+    AnyDtype,
     AnyName,
     AnyPandas,
     DataClass,
@@ -43,7 +43,7 @@ class Field:
     type: Optional[Any]
     """Type (hint) of the field data."""
 
-    dtype: Optional[AnyDType]
+    dtype: Optional[AnyDtype]
     """Data type of the field data."""
 
     default: Any
diff --git a/pandas_dataclasses/core/typing.py b/pandas_dataclasses/core/typing.py
index 02d3a34..bb87075 100644
--- a/pandas_dataclasses/core/typing.py
+++ b/pandas_dataclasses/core/typing.py
@@ -36,7 +36,7 @@ from typing_extensions import (
 
 
 # type hints (private)
-AnyDType: TypeAlias = "np.dtype[Any] | ExtensionDtype"
+AnyDtype: TypeAlias = "np.dtype[Any] | ExtensionDtype"
 AnyName: TypeAlias = "Hashable | dict[Hashable, Hashable]"
 AnyPandas: TypeAlias = "pd.DataFrame | pd.Series"
 P = ParamSpec("P")
@@ -138,7 +138,7 @@ def get_annotations(tp: Any) -> Tuple[Any, ...]:
     raise TypeError("Could not find any role-annotated type.")
 
 
-def get_dtype(tp: Any) -> Optional[AnyDType]:
+def get_dtype(tp: Any) -> Optional[AnyDtype]:
     """Extract a NumPy or pandas data type."""
     try:
         dtype = get_args(get_annotated(tp))[0]


Test patch:
diff --git a/tests/test_parsers.py b/tests/test_parsers.py
index 48cd575..b8a323b 100644
--- a/tests/test_parsers.py
+++ b/tests/test_parsers.py
@@ -87,10 +87,10 @@ def test_get_data() -> None:
     assert keys[2] == spec.fields.of_data[2].hashable_name
     assert keys[3] == spec.fields.of_data[3].hashable_name
 
-    assert values[0].dtype == spec.fields.of_data[0].dtype
-    assert values[1].dtype == spec.fields.of_data[1].dtype
-    assert values[2].dtype == spec.fields.of_data[2].dtype
-    assert values[3].dtype == spec.fields.of_data[3].dtype
+    assert values[0].dtype.numpy_dtype == spec.fields.of_data[0].dtype
+    assert values[1].dtype.numpy_dtype == spec.fields.of_data[1].dtype
+    assert values[2].dtype.numpy_dtype == spec.fields.of_data[2].dtype
+    assert values[3].dtype.numpy_dtype == spec.fields.of_data[3].dtype
 
     assert (values[0] == spec.fields.of_data[0].default).all()
     assert (values[1] == spec.fields.of_data[1].default).all()