diff --git a/pandas/core/dtypes/dtypes.py b/pandas/core/dtypes/dtypes.py index d7ba150e3ec9d..17c4c6ba1c701 100644 --- a/pandas/core/dtypes/dtypes.py +++ b/pandas/core/dtypes/dtypes.py @@ -3,7 +3,18 @@ """ import re -from typing import Any, Dict, List, MutableMapping, Optional, Tuple, Type, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + Dict, + List, + MutableMapping, + Optional, + Tuple, + Type, + Union, + cast, +) import numpy as np import pytz @@ -16,6 +27,15 @@ from pandas.core.dtypes.generic import ABCCategoricalIndex, ABCDateOffset, ABCIndexClass from pandas.core.dtypes.inference import is_bool, is_list_like +if TYPE_CHECKING: + import pyarrow # noqa: F401 + from pandas.core.arrays import ( # noqa: F401 + IntervalArray, + PeriodArray, + DatetimeArray, + ) + from pandas import Categorical # noqa: F401 + str_type = str @@ -68,7 +88,7 @@ def register(self, dtype: Type[ExtensionDtype]) -> None: """ Parameters ---------- - dtype : ExtensionDtype + dtype : ExtensionDtype class """ if not issubclass(dtype, ExtensionDtype): raise ValueError("can only register pandas extension dtypes") @@ -122,7 +142,7 @@ class PandasExtensionDtype(ExtensionDtype): # and ExtensionDtype's @properties in the subclasses below. The kind and # type variables in those subclasses are explicitly typed below. subdtype = None - str: Optional[str_type] = None + str: str_type num = 100 shape: Tuple[int, ...] = tuple() itemsize = 8 @@ -500,7 +520,7 @@ def _hash_categories(categories, ordered: Ordered = True) -> int: return np.bitwise_xor.reduce(hashed) @classmethod - def construct_array_type(cls): + def construct_array_type(cls) -> Type["Categorical"]: """ Return the array type associated with this dtype. @@ -508,7 +528,7 @@ def construct_array_type(cls): ------- type """ - from pandas import Categorical + from pandas import Categorical # noqa: F811 return Categorical @@ -672,9 +692,9 @@ class DatetimeTZDtype(PandasExtensionDtype): _match = re.compile(r"(datetime64|M8)\[(?P.+), (?P.+)\]") _cache: Dict[str_type, PandasExtensionDtype] = {} - def __init__(self, unit="ns", tz=None): + def __init__(self, unit: Union[str_type, "DatetimeTZDtype"] = "ns", tz=None): if isinstance(unit, DatetimeTZDtype): - unit, tz = unit.unit, unit.tz + unit, tz = unit.unit, unit.tz # type: ignore if unit != "ns": if isinstance(unit, str) and tz is None: @@ -704,7 +724,7 @@ def __init__(self, unit="ns", tz=None): self._tz = tz @property - def unit(self): + def unit(self) -> str_type: """ The precision of the datetime data. """ @@ -718,7 +738,7 @@ def tz(self): return self._tz @classmethod - def construct_array_type(cls): + def construct_array_type(cls) -> Type["DatetimeArray"]: """ Return the array type associated with this dtype. @@ -726,12 +746,12 @@ def construct_array_type(cls): ------- type """ - from pandas.core.arrays import DatetimeArray + from pandas.core.arrays import DatetimeArray # noqa: F811 return DatetimeArray @classmethod - def construct_from_string(cls, string: str_type): + def construct_from_string(cls, string: str_type) -> "DatetimeTZDtype": """ Construct a DatetimeTZDtype from a string. @@ -789,7 +809,7 @@ def __eq__(self, other: Any) -> bool: and str(self.tz) == str(other.tz) ) - def __setstate__(self, state): + def __setstate__(self, state) -> None: # for pickle compat. __get_state__ is defined in the # PandasExtensionDtype superclass and uses the public properties to # pickle -> need to set the settable private ones here (see GH26067) @@ -884,7 +904,7 @@ def _parse_dtype_strict(cls, freq): raise ValueError("could not construct PeriodDtype") @classmethod - def construct_from_string(cls, string): + def construct_from_string(cls, string: str_type) -> "PeriodDtype": """ Strict construction from a string, raise a TypeError if not possible @@ -934,7 +954,7 @@ def __setstate__(self, state): self._freq = state["freq"] @classmethod - def is_dtype(cls, dtype) -> bool: + def is_dtype(cls, dtype: object) -> bool: """ Return a boolean if we if the passed type is an actual dtype that we can match (via string or type) @@ -955,7 +975,7 @@ def is_dtype(cls, dtype) -> bool: return super().is_dtype(dtype) @classmethod - def construct_array_type(cls): + def construct_array_type(cls) -> Type["PeriodArray"]: """ Return the array type associated with this dtype. @@ -967,9 +987,13 @@ def construct_array_type(cls): return PeriodArray - def __from_arrow__(self, array): - """Construct PeriodArray from pyarrow Array/ChunkedArray.""" - import pyarrow + def __from_arrow__( + self, array: Union["pyarrow.Array", "pyarrow.ChunkedArray"] + ) -> "PeriodArray": + """ + Construct PeriodArray from pyarrow Array/ChunkedArray. + """ + import pyarrow # noqa: F811 from pandas.core.arrays import PeriodArray from pandas.core.arrays._arrow_utils import pyarrow_array_to_numpy_and_mask @@ -1075,7 +1099,7 @@ def subtype(self): return self._subtype @classmethod - def construct_array_type(cls): + def construct_array_type(cls) -> Type["IntervalArray"]: """ Return the array type associated with this dtype. @@ -1142,7 +1166,7 @@ def __setstate__(self, state): self._subtype = state["subtype"] @classmethod - def is_dtype(cls, dtype) -> bool: + def is_dtype(cls, dtype: object) -> bool: """ Return a boolean if we if the passed type is an actual dtype that we can match (via string or type) @@ -1160,9 +1184,13 @@ def is_dtype(cls, dtype) -> bool: return False return super().is_dtype(dtype) - def __from_arrow__(self, array): - """Construct IntervalArray from pyarrow Array/ChunkedArray.""" - import pyarrow + def __from_arrow__( + self, array: Union["pyarrow.Array", "pyarrow.ChunkedArray"] + ) -> "IntervalArray": + """ + Construct IntervalArray from pyarrow Array/ChunkedArray. + """ + import pyarrow # noqa: F811 from pandas.core.arrays import IntervalArray if isinstance(array, pyarrow.Array):