test dtype

Author: cmp0xff

pandas-stubs/core/construction.pyi

@@ -40,7 +40,6 @@ 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 (
-    BuiltinBooleanDtypeArg,
     BuiltinFloatDtypeArg,
     BuiltinIntDtypeArg,
     BuiltinStrDtypeArg,
@@ -114,7 +113,7 @@ def array(
 @overload
 def array(  # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
     data: Sequence[bool | np.bool | NAType | None] | np_ndarray_bool | BooleanArray,
-    dtype: BuiltinBooleanDtypeArg | PandasBooleanDtypeArg | None = None,
+    dtype: PandasBooleanDtypeArg | None = None,
     copy: bool = True,
 ) -> BooleanArray: ...
 @overload

tests/__init__.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from collections.abc import Generator
 from contextlib import (
     AbstractContextManager,
     nullcontext,
@@ -31,13 +32,15 @@
 if TYPE_CHECKING:
     from pandas._typing import (
         BooleanDtypeArg as BooleanDtypeArg,
+        BuiltinBooleanDtypeArg as BuiltinBooleanDtypeArg,
         BytesDtypeArg as BytesDtypeArg,
         CategoryDtypeArg as CategoryDtypeArg,
         ComplexDtypeArg as ComplexDtypeArg,
         Dtype as Dtype,
         FloatDtypeArg as FloatDtypeArg,
         IntDtypeArg as IntDtypeArg,
         ObjectDtypeArg as ObjectDtypeArg,
+        PandasBooleanDtypeArg as PandasBooleanDtypeArg,
         StrDtypeArg as StrDtypeArg,
         T as T,
         TimedeltaDtypeArg as TimedeltaDtypeArg,
@@ -67,6 +70,10 @@
         np_ndarray_td as np_ndarray_td,
     )
 else:
+    # Builtin bool type and its string alias
+    BuiltinBooleanDtypeArg: TypeAlias = type[bool] | Literal["bool"]
+    # Pandas nullable boolean type and its string alias
+    PandasBooleanDtypeArg: TypeAlias = pd.BooleanDtype | Literal["boolean"]
     _G = TypeVar("_G", bound=np.generic)
     _S = TypeVar("_S", bound=tuple[int, ...])
     # Separately define here so pytest works
@@ -244,3 +251,14 @@ def pytest_warns_bounded(
     if upper_exception is None:
         return nullcontext()
     return suppress(upper_exception)
+
+
+def get_dtype(dtype: object) -> Generator[type | str, None, None]:
+    """Extract types and string literals from a Union or Literal type."""
+    if isinstance(dtype, str):
+        yield dtype
+    elif isinstance(dtype, type):
+        yield dtype()
+    else:
+        for arg in get_args(dtype):
+            yield from get_dtype(arg)

tests/arrays/test_boolean_array.py

@@ -1,9 +1,14 @@
 import numpy as np
 import pandas as pd
 from pandas.core.arrays.boolean import BooleanArray
+import pytest
 from typing_extensions import assert_type
 
-from tests import check
+from tests import (
+    PandasBooleanDtypeArg,
+    check,
+    get_dtype,
+)
 
 
 def test_constructor() -> None:
@@ -15,3 +20,10 @@ def test_constructor() -> None:
     check(assert_type(pd.array(np.array([1], np.bool_)), BooleanArray), BooleanArray)
 
     check(assert_type(pd.array(pd.array([True])), BooleanArray), BooleanArray)
+
+    pd.array([True], dtype=pd.BooleanDtype())
+
+
+@pytest.mark.parametrize("dtype", get_dtype(PandasBooleanDtypeArg))
+def test_constructors_dtype(dtype: PandasBooleanDtypeArg):
+    check(assert_type(pd.array([True], dtype=dtype), BooleanArray), BooleanArray)