Add type hints to dtypes/dtypes.py (CategoricalDtype)

Author: topper-123

pandas/core/arrays/categorical.py

@@ -429,7 +429,7 @@ def ordered(self):
         return self.dtype.ordered
 
     @property
-    def dtype(self):
+    def dtype(self) -> 'CategoricalDtype':
         """
         The :class:`~pandas.api.types.CategoricalDtype` for this instance
         """

pandas/core/dtypes/dtypes.py

@@ -1,6 +1,7 @@
 """ define extension dtypes """
 import re
-from typing import Any, Dict, Optional, Tuple, Type
+import typing
+from typing import Any, Dict, List, Optional, Tuple, Type
 import warnings
 
 import numpy as np
@@ -15,10 +16,8 @@
 from .base import ExtensionDtype, _DtypeOpsMixin
 from .inference import is_list_like
 
-str_type = str
 
-
-def register_extension_dtype(cls):
+def register_extension_dtype(cls: 'ExtensionDtype') -> 'ExtensionDtype':
     """
     Register an ExtensionType with pandas as class decorator.
 
@@ -60,20 +59,20 @@ class Registry:
     These are tried in order.
     """
     def __init__(self):
-        self.dtypes = []
+        self.dtypes = []  # type: List[ExtensionDtype]
 
-    def register(self, dtype):
+    def register(self, dtype: 'ExtensionDtype') -> None:
         """
         Parameters
         ----------
         dtype : ExtensionDtype
         """
-        if not issubclass(dtype, (PandasExtensionDtype, ExtensionDtype)):
+        if not issubclass(dtype, ExtensionDtype):
             raise ValueError("can only register pandas extension dtypes")
 
         self.dtypes.append(dtype)
 
-    def find(self, dtype):
+    def find(self, dtype: 'ExtensionDtype') -> Optional[ExtensionDtype]:
         """
         Parameters
         ----------
@@ -117,25 +116,25 @@ class PandasExtensionDtype(_DtypeOpsMixin):
     # and ExtensionDtype's @properties in the subclasses below. The kind and
     # type variables in those subclasses are explicitly typed below.
     subdtype = None
-    str = None  # type: Optional[str_type]
+    str = None  # type: Optional[str]
     num = 100
     shape = tuple()  # type: Tuple[int, ...]
     itemsize = 8
     base = None
     isbuiltin = 0
     isnative = 0
-    _cache = {}  # type: Dict[str_type, 'PandasExtensionDtype']
+    _cache = {}  # type: Dict[str, 'PandasExtensionDtype']
 
-    def __unicode__(self):
+    def __unicode__(self) -> str:
         return self.name
 
-    def __str__(self):
+    def __str__(self) -> str:
         """
         Return a string representation for a particular Object
         """
         return self.__unicode__()
 
-    def __bytes__(self):
+    def __bytes__(self) -> bytes:
         """
         Return a string representation for a particular object.
         """
@@ -144,22 +143,22 @@ def __bytes__(self):
         encoding = get_option("display.encoding")
         return self.__unicode__().encode(encoding, 'replace')
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """
         Return a string representation for a particular object.
         """
         return str(self)
 
-    def __hash__(self):
+    def __hash__(self) -> int:
         raise NotImplementedError("sub-classes should implement an __hash__ "
                                   "method")
 
-    def __getstate__(self):
+    def __getstate__(self) -> Dict[str, Any]:
         # pickle support; we don't want to pickle the cache
         return {k: getattr(self, k, None) for k in self._metadata}
 
     @classmethod
-    def reset_cache(cls):
+    def reset_cache(cls) -> None:
         """ clear the cache """
         cls._cache = {}
 
@@ -217,23 +216,30 @@ class CategoricalDtype(PandasExtensionDtype, ExtensionDtype):
     # TODO: Document public vs. private API
     name = 'category'
     type = CategoricalDtypeType  # type: Type[CategoricalDtypeType]
-    kind = 'O'  # type: str_type
+    kind = 'O'  # type: str
     str = '|O08'
     base = np.dtype('O')
     _metadata = ('categories', 'ordered')
-    _cache = {}  # type: Dict[str_type, PandasExtensionDtype]
+    _cache = {}  # type: Dict[str, PandasExtensionDtype]
 
-    def __init__(self, categories=None, ordered=None):
+    def __init__(self, categories=None, ordered: bool = None):
         self._finalize(categories, ordered, fastpath=False)
 
     @classmethod
-    def _from_fastpath(cls, categories=None, ordered=None):
+    def _from_fastpath(cls,
+                       categories=None,
+                       ordered: bool = None
+                       ) -> 'CategoricalDtype':
         self = cls.__new__(cls)
         self._finalize(categories, ordered, fastpath=True)
         return self
 
     @classmethod
-    def _from_categorical_dtype(cls, dtype, categories=None, ordered=None):
+    def _from_categorical_dtype(cls,
+                                dtype: 'CategoricalDtype',
+                                categories=None,
+                                ordered: bool = None,
+                                ) -> 'CategoricalDtype':
         if categories is ordered is None:
             return dtype
         if categories is None:
@@ -243,8 +249,12 @@ def _from_categorical_dtype(cls, dtype, categories=None, ordered=None):
         return cls(categories, ordered)
 
     @classmethod
-    def _from_values_or_dtype(cls, values=None, categories=None, ordered=None,
-                              dtype=None):
+    def _from_values_or_dtype(cls,
+                              values=None,
+                              categories=None,
+                              ordered: bool = None,
+                              dtype: 'CategoricalDtype' = None,
+                              ) -> 'CategoricalDtype':
         """
         Construct dtype from the input parameters used in :class:`Categorical`.
 
@@ -326,7 +336,11 @@ def _from_values_or_dtype(cls, values=None, categories=None, ordered=None,
 
         return dtype
 
-    def _finalize(self, categories, ordered, fastpath=False):
+    def _finalize(self,
+                  categories,
+                  ordered: Optional[bool],
+                  fastpath: bool = False,
+                  ) -> None:
 
         if ordered is not None:
             self.validate_ordered(ordered)
@@ -338,14 +352,14 @@ def _finalize(self, categories, ordered, fastpath=False):
         self._categories = categories
         self._ordered = ordered
 
-    def __setstate__(self, state):
+    def __setstate__(self, state: 'Dict[str, Any]') -> 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)
         self._categories = state.pop('categories', None)
         self._ordered = state.pop('ordered', False)
 
-    def __hash__(self):
+    def __hash__(self) -> int:
         # _hash_categories returns a uint64, so use the negative
         # space for when we have unknown categories to avoid a conflict
         if self.categories is None:
@@ -356,7 +370,7 @@ def __hash__(self):
         # We *do* want to include the real self.ordered here
         return int(self._hash_categories(self.categories, self.ordered))
 
-    def __eq__(self, other):
+    def __eq__(self, other: Any) -> bool:
         """
         Rules for CDT equality:
         1) Any CDT is equal to the string 'category'
@@ -403,7 +417,7 @@ def __repr__(self):
         return tpl.format(data, self.ordered)
 
     @staticmethod
-    def _hash_categories(categories, ordered=True):
+    def _hash_categories(categories, ordered: bool = True) -> int:
         from pandas.core.util.hashing import (
             hash_array, _combine_hash_arrays, hash_tuples
         )
@@ -453,7 +467,7 @@ def construct_array_type(cls):
         return Categorical
 
     @classmethod
-    def construct_from_string(cls, string):
+    def construct_from_string(cls, string: str) -> 'CategoricalDtype':
         """
         attempt to construct this type from a string, raise a TypeError if
         it's not possible """
@@ -466,7 +480,7 @@ def construct_from_string(cls, string):
             pass
 
     @staticmethod
-    def validate_ordered(ordered):
+    def validate_ordered(ordered: bool) -> None:
         """
         Validates that we have a valid ordered parameter. If
         it is not a boolean, a TypeError will be raised.
@@ -486,7 +500,7 @@ def validate_ordered(ordered):
             raise TypeError("'ordered' must either be 'True' or 'False'")
 
     @staticmethod
-    def validate_categories(categories, fastpath=False):
+    def validate_categories(categories, fastpath: bool = False):
         """
         Validates that we have good categories
 
@@ -519,9 +533,9 @@ def validate_categories(categories, fastpath=False):
         if isinstance(categories, ABCCategoricalIndex):
             categories = categories.categories
 
-        return categories
+        return typing.cast(Index, categories)
 
-    def update_dtype(self, dtype):
+    def update_dtype(self, dtype: 'CategoricalDtype') -> 'CategoricalDtype':
         """
         Returns a CategoricalDtype with categories and ordered taken from dtype
         if specified, otherwise falling back to self if unspecified
@@ -560,17 +574,18 @@ def categories(self):
         """
         An ``Index`` containing the unique categories allowed.
         """
-        return self._categories
+        from pandas import Index
+        return typing.cast(Index, self._categories)
 
     @property
-    def ordered(self):
+    def ordered(self) -> bool:
         """
         Whether the categories have an ordered relationship.
         """
         return self._ordered
 
     @property
-    def _is_boolean(self):
+    def _is_boolean(self) -> bool:
         from pandas.core.dtypes.common import is_bool_dtype
 
         return is_bool_dtype(self.categories)
@@ -614,14 +629,14 @@ class DatetimeTZDtype(PandasExtensionDtype, ExtensionDtype):
     datetime64[ns, tzfile('/usr/share/zoneinfo/US/Central')]
     """
     type = Timestamp  # type: Type[Timestamp]
-    kind = 'M'  # type: str_type
+    kind = 'M'  # type: str
     str = '|M8[ns]'
     num = 101
     base = np.dtype('M8[ns]')
     na_value = NaT
     _metadata = ('unit', 'tz')
     _match = re.compile(r"(datetime64|M8)\[(?P<unit>.+), (?P<tz>.+)\]")
-    _cache = {}  # type: Dict[str_type, PandasExtensionDtype]
+    _cache = {}  # type: Dict[str, PandasExtensionDtype]
 
     def __init__(self, unit="ns", tz=None):
         if isinstance(unit, DatetimeTZDtype):
@@ -765,13 +780,13 @@ class PeriodDtype(ExtensionDtype, PandasExtensionDtype):
     period[M]
     """
     type = Period  # type: Type[Period]
-    kind = 'O'  # type: str_type
+    kind = 'O'  # type: str
     str = '|O08'
     base = np.dtype('O')
     num = 102
     _metadata = ('freq',)
     _match = re.compile(r"(P|p)eriod\[(?P<freq>.+)\]")
-    _cache = {}  # type: Dict[str_type, PandasExtensionDtype]
+    _cache = {}  # type: Dict[str, PandasExtensionDtype]
 
     def __new__(cls, freq=None):
         """
@@ -919,13 +934,13 @@ class IntervalDtype(PandasExtensionDtype, ExtensionDtype):
     interval[int64]
     """
     name = 'interval'
-    kind = None  # type: Optional[str_type]
+    kind = None  # type: Optional[str]
     str = '|O08'
     base = np.dtype('O')
     num = 103
     _metadata = ('subtype',)
     _match = re.compile(r"(I|i)nterval\[(?P<subtype>.+)\]")
-    _cache = {}  # type: Dict[str_type, PandasExtensionDtype]
+    _cache = {}  # type: Dict[str, PandasExtensionDtype]
 
     def __new__(cls, subtype=None):
         from pandas.core.dtypes.common import (