Add NpDtype for functions only allow numpy dtypes

Author: avinashpancham

pandas/_typing.py

@@ -94,9 +94,8 @@
 Axes = Collection
 
 # dtypes
-Dtype = Union[
-    "ExtensionDtype", str, np.dtype, Type[Union[str, float, int, complex, bool, object]]
-]
+NpDtype = Union[str, np.dtype, Type[Union[str, float, int, complex, bool, object]]]
+Dtype = Union["ExtensionDtype", NpDtype]
 # DtypeArg specifies all allowable dtypes in a functions its dtype argument
 DtypeArg = Union[Dtype, Dict[Label, Dtype]]
 DtypeObj = Union[np.dtype, "ExtensionDtype"]

pandas/core/arrays/numpy_.py

@@ -5,7 +5,7 @@
 from numpy.lib.mixins import NDArrayOperatorsMixin
 
 from pandas._libs import lib
-from pandas._typing import Dtype, Scalar
+from pandas._typing import Dtype, NpDtype, Scalar
 from pandas.compat.numpy import function as nv
 
 from pandas.core.dtypes.dtypes import ExtensionDtype
@@ -199,7 +199,7 @@ def dtype(self) -> PandasDtype:
     # ------------------------------------------------------------------------
     # NumPy Array Interface
 
-    def __array__(self, dtype=None) -> np.ndarray:
+    def __array__(self, dtype: Optional[NpDtype] = None) -> np.ndarray:
         return np.asarray(self._ndarray, dtype=dtype)
 
     _HANDLED_TYPES = (np.ndarray, numbers.Number)

pandas/core/arrays/period.py

@@ -26,7 +26,7 @@
     get_period_field_arr,
     period_asfreq_arr,
 )
-from pandas._typing import AnyArrayLike, Dtype
+from pandas._typing import AnyArrayLike, Dtype, NpDtype
 from pandas.util._decorators import cache_readonly, doc
 
 from pandas.core.dtypes.common import (
@@ -304,7 +304,7 @@ def freq(self) -> BaseOffset:
         """
         return self.dtype.freq
 
-    def __array__(self, dtype=None) -> np.ndarray:
+    def __array__(self, dtype: Optional[NpDtype] = None) -> np.ndarray:
         if dtype == "i8":
             return self.asi8
         elif dtype == bool:

pandas/core/indexes/base.py

@@ -27,7 +27,16 @@
 from pandas._libs.lib import is_datetime_array, no_default
 from pandas._libs.tslibs import IncompatibleFrequency, OutOfBoundsDatetime, Timestamp
 from pandas._libs.tslibs.timezones import tz_compare
-from pandas._typing import AnyArrayLike, ArrayLike, Dtype, DtypeObj, Label, Shape, final
+from pandas._typing import (
+    AnyArrayLike,
+    ArrayLike,
+    Dtype,
+    DtypeObj,
+    Label,
+    NpDtype,
+    Shape,
+    final,
+)
 from pandas.compat.numpy import function as nv
 from pandas.errors import DuplicateLabelError, InvalidIndexError
 from pandas.util._decorators import Appender, cache_readonly, doc
@@ -619,7 +628,7 @@ def __len__(self) -> int:
         """
         return len(self._data)
 
-    def __array__(self, dtype=None) -> np.ndarray:
+    def __array__(self, dtype: Optional[NpDtype] = None) -> np.ndarray:
         """
         The array interface, return my values.
         """

pandas/core/indexes/multi.py

@@ -20,7 +20,7 @@
 
 from pandas._libs import algos as libalgos, index as libindex, lib
 from pandas._libs.hashtable import duplicated_int64
-from pandas._typing import AnyArrayLike, Dtype, DtypeObj, Label, Scalar, Shape
+from pandas._typing import AnyArrayLike, Dtype, DtypeObj, Label, NpDtype, Scalar, Shape
 from pandas.compat.numpy import function as nv
 from pandas.errors import InvalidIndexError, PerformanceWarning, UnsortedIndexError
 from pandas.util._decorators import Appender, cache_readonly, doc
@@ -1204,7 +1204,7 @@ def copy(
             new_index = new_index.astype(dtype)
         return new_index
 
-    def __array__(self, dtype=None) -> np.ndarray:
+    def __array__(self, dtype: Optional[NpDtype] = None) -> np.ndarray:
         """ the array interface, return my values """
         return self.values