interval array

Author: cmp0xff

pandas-stubs/core/construction.pyi

@@ -49,6 +49,7 @@ from pandas._typing import (
     StrDtypeArg,
     TimedeltaDtypeArg,
     TimestampDtypeArg,
+    UIntDtypeArg,
     np_ndarray,
     np_ndarray_anyint,
     np_ndarray_bool,
@@ -91,7 +92,7 @@ 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: IntDtypeArg | None = None,
+    dtype: IntDtypeArg | UIntDtypeArg | None = None,
     copy: bool = True,
 ) -> IntegerArray: ...
 @overload
@@ -143,12 +144,8 @@ def array(  # type: ignore[overload-overlap]
 ) -> PeriodArray: ...
 @overload
 def array(  # type: ignore[overload-overlap]
-    data: (
-        Sequence[IntervalT | NAType | None]
-        | IntervalIndex
-        | Series[Interval]
-        | IntervalArray
-    ),
+    # float("nan") also works, but I don't know how to put it in
+    data: Sequence[IntervalT | None] | IntervalArray | IntervalIndex | Series[Interval],
     dtype: IntervalDtype | None = None,
     copy: bool = True,
 ) -> IntervalArray: ...

tests/arrays/test_interval_array.py

@@ -0,0 +1,32 @@
+import pandas as pd
+from pandas.core.arrays.interval import IntervalArray
+from typing_extensions import assert_type
+
+from tests import check
+
+
+def test_constructor() -> None:
+    itv = pd.Interval(0, 1)
+    check(
+        assert_type(  # type: ignore[assert-type]  # I do not understand
+            pd.array([itv]), IntervalArray
+        ),
+        IntervalArray,
+    )
+    check(
+        assert_type(  # type: ignore[assert-type]  # I do not understand
+            pd.array([itv, None]), IntervalArray
+        ),
+        IntervalArray,
+    )
+
+    check(
+        assert_type(  # type: ignore[assert-type]  # I do not understand
+            pd.array(pd.array([itv])), IntervalArray
+        ),
+        IntervalArray,
+    )
+
+    check(assert_type(pd.array(pd.Index([itv])), IntervalArray), IntervalArray)
+
+    check(assert_type(pd.array(pd.Series([itv])), IntervalArray), IntervalArray)