dtype numpy extension array

Author: cmp0xff

pandas-stubs/core/construction.pyi

@@ -40,12 +40,12 @@ from pandas._libs.tslibs.period import Period
 from pandas._libs.tslibs.timedeltas import Timedelta
 from pandas._libs.tslibs.timestamps import Timestamp
 from pandas._typing import (
-    BuiltinFloatDtypeArg,
-    BuiltinIntDtypeArg,
-    BuiltinStrDtypeArg,
+    BuiltinDtypeArg,
     CategoryDtypeArg,
     IntervalT,
-    NumpyDtypeArg,
+    NumpyNotTimeDtypeArg,
+    NumpyTimedeltaDtypeArg,
+    NumpyTimestampDtypeArg,
     PandasBooleanDtypeArg,
     PandasFloatDtypeArg,
     PandasIntDtypeArg,
@@ -125,15 +125,15 @@ def array(  # type: ignore[overload-overlap]
 @overload
 def array(  # type: ignore[overload-overlap]
     data: Sequence[int | np.integer | NAType | None] | np_ndarray_anyint | IntegerArray,
-    dtype: BuiltinIntDtypeArg | PandasIntDtypeArg | PandasUIntDtypeArg | None = None,
+    dtype: PandasIntDtypeArg | PandasUIntDtypeArg | None = None,
     copy: bool = True,
 ) -> IntegerArray: ...
 @overload
 def array(  # type: ignore[overload-overlap]
     data: (
         Sequence[float | np.floating | NAType | None] | np_ndarray_float | FloatingArray
     ),
-    dtype: BuiltinFloatDtypeArg | PandasFloatDtypeArg | None = None,
+    dtype: PandasFloatDtypeArg | None = None,
     copy: bool = True,
 ) -> FloatingArray: ...
 @overload
@@ -146,7 +146,7 @@ def array(  # type: ignore[overload-overlap]
         | DatetimeIndex
         | Series[Timestamp]
     ),
-    dtype: PandasTimestampDtypeArg | None = None,
+    dtype: PandasTimestampDtypeArg | NumpyTimestampDtypeArg | None = None,
     copy: bool = True,
 ) -> DatetimeArray: ...
 @overload
@@ -158,13 +158,13 @@ def array(
         | TimedeltaIndex
         | Series[Timedelta]
     ),
-    dtype: None = None,
+    dtype: NumpyTimedeltaDtypeArg | None = None,
     copy: bool = True,
 ) -> TimedeltaArray: ...
 @overload
 def array(  # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
     data: SequenceNotStr[str | np.str_ | NAType | None] | np_ndarray_str | StringArray,
-    dtype: BuiltinStrDtypeArg | PandasStrDtypeArg | None = None,
+    dtype: PandasStrDtypeArg | None = None,
     copy: bool = True,
 ) -> StringArray: ...
 @overload
@@ -181,7 +181,7 @@ def array(  # type: ignore[overload-overlap]
 @overload
 def array(
     data: SequenceNotStr[object] | np_ndarray | NumpyExtensionArray | RangeIndex,
-    dtype: NumpyDtypeArg | None = None,
+    dtype: BuiltinDtypeArg | NumpyNotTimeDtypeArg | None = None,
     copy: bool = True,
 ) -> NumpyExtensionArray: ...
 @overload

tests/_typing.py

@@ -8,11 +8,13 @@
 
 from pandas._typing import (
     BooleanDtypeArg,
+    BuiltinDtypeArg,
     BytesDtypeArg,
     CategoryDtypeArg,
     ComplexDtypeArg,
     FloatDtypeArg,
     IntDtypeArg,
+    NumpyNotTimeDtypeArg,
     ObjectDtypeArg,
     PandasAstypeComplexDtypeArg,
     PandasAstypeFloatDtypeArg,
@@ -62,6 +64,8 @@
     "ComplexDtypeArg",
     "IntDtypeArg",
     "np_ndarray_str",
+    "BuiltinDtypeArg",
+    "NumpyNotTimeDtypeArg",
     "ObjectDtypeArg",
     "PandasAstypeComplexDtypeArg",
     "PandasAstypeTimedeltaDtypeArg",

tests/arrays/test_datetime_array.py

@@ -24,12 +24,8 @@ def test_constructor() -> None:
     dt_nat = cast(list[np.datetime64 | NaTType], [np_dt, pd.NaT])
     check(assert_type(pd.array(dt_nat), DatetimeArray), DatetimeArray)
 
-    check(
-        assert_type(  # type: ignore[assert-type] # I do not understand
-            pd.array(np.array([dt], np.datetime64)), DatetimeArray
-        ),
-        DatetimeArray,
-    )
+    np_arr = np.array([dt], np.datetime64)
+    check(assert_type(pd.array(np_arr), DatetimeArray), DatetimeArray)
 
     check(assert_type(pd.array(pd.array([dt])), DatetimeArray), DatetimeArray)
 

tests/arrays/test_integer_array.py

@@ -12,11 +12,7 @@ def test_constructor() -> None:
     check(assert_type(pd.array([1, None]), IntegerArray), IntegerArray)
     check(assert_type(pd.array([1, pd.NA, None]), IntegerArray), IntegerArray)
 
-    check(
-        assert_type(  # type: ignore[assert-type] # I do not understand
-            pd.array(np.array([1], np.int64)), IntegerArray
-        ),
-        IntegerArray,
-    )
+    np_arr = np.array([1], np.int64)
+    check(assert_type(pd.array(np_arr), IntegerArray), IntegerArray)
 
     check(assert_type(pd.array(pd.array([1])), IntegerArray), IntegerArray)

tests/arrays/test_numpy_extension_array.py

@@ -1,28 +1,34 @@
 import numpy as np
 import pandas as pd
 from pandas.core.arrays.numpy_ import NumpyExtensionArray
+import pytest
 from typing_extensions import assert_type
 
-from tests import check
+from tests import (
+    check,
+    get_dtype,
+)
+from tests._typing import (
+    BuiltinDtypeArg,
+    NumpyNotTimeDtypeArg,
+)
 
 
 def test_constructor() -> None:
     # check(
     #     assert_type(pd.array([pd.NA, None]), NumpyExtensionArray), NumpyExtensionArray
     # )
 
+    pd_arr = pd.array([1, "🐼"])
     check(
         assert_type(  # type: ignore[assert-type] # I do not understand
-            pd.array([1, "🐼"]), NumpyExtensionArray
-        ),
-        NumpyExtensionArray,
-    )
-    check(
-        assert_type(  # type: ignore[assert-type] # I do not understand
-            pd.array(np.array([1, "🐼"], np.object_)), NumpyExtensionArray
+            pd_arr, NumpyExtensionArray
         ),
         NumpyExtensionArray,
     )
+
+    np_arr = np.array([1, "🐼"], np.object_)
+    check(assert_type(pd.array(np_arr), NumpyExtensionArray), NumpyExtensionArray)
     # check(
     #     assert_type(pd.array(pd.array([pd.NA, None])), NumpyExtensionArray),
     #     NumpyExtensionArray,
@@ -31,3 +37,17 @@ def test_constructor() -> None:
         assert_type(pd.array(pd.RangeIndex(0, 1)), NumpyExtensionArray),
         NumpyExtensionArray,
     )
+
+
+@pytest.mark.parametrize("dtype", get_dtype(BuiltinDtypeArg | NumpyNotTimeDtypeArg))
+def test_constructors_dtype(dtype: BuiltinDtypeArg | NumpyNotTimeDtypeArg):
+    if dtype == "V" or "void" in str(dtype):
+        check(
+            assert_type(pd.array([b"1"], dtype=dtype), NumpyExtensionArray),
+            NumpyExtensionArray,
+        )
+    else:
+        check(
+            assert_type(pd.array([1], dtype=dtype), NumpyExtensionArray),
+            NumpyExtensionArray,
+        )

tests/arrays/test_timedelta_array.py

@@ -19,12 +19,8 @@ def test_constructor() -> None:
     check(assert_type(pd.array([td, None]), TimedeltaArray), TimedeltaArray)
     check(assert_type(pd.array([td, pd.NaT, None]), TimedeltaArray), TimedeltaArray)
 
-    check(
-        assert_type(  # type: ignore[assert-type] # I do not understand
-            pd.array(np.array([td], np.timedelta64)), TimedeltaArray
-        ),
-        TimedeltaArray,
-    )
+    np_arr = np.array([td], np.timedelta64)
+    check(assert_type(pd.array(np_arr), TimedeltaArray), TimedeltaArray)
 
     check(assert_type(pd.array(pd.array([td])), TimedeltaArray), TimedeltaArray)