diff --git a/pandas/core/arrays/categorical.py b/pandas/core/arrays/categorical.py index 155638aca5560..9f5e3e8ee77f0 100644 --- a/pandas/core/arrays/categorical.py +++ b/pandas/core/arrays/categorical.py @@ -435,7 +435,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 """ diff --git a/pandas/core/dtypes/dtypes.py b/pandas/core/dtypes/dtypes.py index 7fe8ce7d71683..9da6fb84ee18b 100644 --- a/pandas/core/dtypes/dtypes.py +++ b/pandas/core/dtypes/dtypes.py @@ -1,6 +1,6 @@ """ define extension dtypes """ import re -from typing import Any, Dict, Optional, Tuple, Type +from typing import Any, Dict, List, Optional, Tuple, Type, Union import warnings import numpy as np @@ -18,7 +18,8 @@ str_type = str -def register_extension_dtype(cls): +def register_extension_dtype(cls: Type[ExtensionDtype], + ) -> Type[ExtensionDtype]: """ Register an ExtensionType with pandas as class decorator. @@ -60,9 +61,9 @@ class Registry: These are tried in order. """ def __init__(self): - self.dtypes = [] + self.dtypes = [] # type: List[Type[ExtensionDtype]] - def register(self, dtype): + def register(self, dtype: Type[ExtensionDtype]) -> None: """ Parameters ---------- @@ -73,11 +74,13 @@ def register(self, dtype): self.dtypes.append(dtype) - def find(self, dtype): + def find(self, + dtype: Union[Type[ExtensionDtype], str], + ) -> Optional[Type[ExtensionDtype]]: """ Parameters ---------- - dtype : PandasExtensionDtype or string + dtype : Type[ExtensionDtype] or string Returns ------- @@ -126,28 +129,28 @@ class PandasExtensionDtype(ExtensionDtype): isnative = 0 _cache = {} # type: Dict[str_type, 'PandasExtensionDtype'] - def __str__(self): + def __str__(self) -> str_type: """ Return a string representation for a particular Object """ return self.name - def __repr__(self): + def __repr__(self) -> str_type: """ 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_type, 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 = {} @@ -211,17 +214,24 @@ class CategoricalDtype(PandasExtensionDtype, ExtensionDtype): _metadata = ('categories', 'ordered') _cache = {} # type: Dict[str_type, 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: @@ -231,8 +241,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`. @@ -314,7 +328,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) @@ -326,14 +344,14 @@ def _finalize(self, categories, ordered, fastpath=False): self._categories = categories self._ordered = ordered - def __setstate__(self, state): + def __setstate__(self, state: Dict[str_type, 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: @@ -344,7 +362,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' @@ -391,7 +409,7 @@ def __repr__(self): return tpl.format(data, self.ordered) @staticmethod - def _hash_categories(categories, ordered=True): + def _hash_categories(categories, ordered: Optional[bool] = True) -> int: from pandas.core.util.hashing import ( hash_array, _combine_hash_arrays, hash_tuples ) @@ -441,7 +459,7 @@ def construct_array_type(cls): return Categorical @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. @@ -461,7 +479,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 @@ -475,7 +493,7 @@ def validate_categories(categories, fastpath=False): ------- categories : Index """ - from pandas import Index + from pandas.core.indexes.base import Index if not fastpath and not is_list_like(categories): msg = "Parameter 'categories' must be list-like, was {!r}" @@ -496,7 +514,7 @@ def validate_categories(categories, fastpath=False): return 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 @@ -538,14 +556,14 @@ def categories(self): return self._categories @property - def ordered(self): + def ordered(self) -> Optional[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)