TYP: pandas/core/dtypes/dtypes.py (#31384)

Author: simonjayhawkins

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,15 +520,15 @@ 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.
 
  Returns
  -------
  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<unit>.+), (?P<tz>.+)\]")
  _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,20 +738,20 @@ 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.
 
  Returns
  -------
  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):