TYPIMG: Partial typing of Categorical

Author: topper-123

pandas/_typing.py

@@ -7,7 +7,7 @@
 from pandas._libs.tslibs.period import Period
 from pandas._libs.tslibs.timedeltas import Timedelta
 
-from pandas.core.dtypes.dtypes import ExtensionDtype
+from pandas.core.dtypes import dtypes
 from pandas.core.dtypes.generic import (
     ABCDataFrame,
     ABCExtensionArray,
@@ -26,8 +26,11 @@
 )
 ArrayLike = TypeVar("ArrayLike", ABCExtensionArray, np.ndarray)
 DatetimeLikeScalar = TypeVar("DatetimeLikeScalar", Period, Timestamp, Timedelta)
-Dtype = Union[str, np.dtype, ExtensionDtype]
+Dtype = Union[str, np.dtype, "dtypes.ExtensionDtype"]
 FilePathOrBuffer = Union[str, Path, IO[AnyStr]]
 
 FrameOrSeries = TypeVar("FrameOrSeries", ABCSeries, ABCDataFrame)
 Scalar = Union[str, int, float]
+
+# TODO(GH26403): Replace with Optional[bool] or bool
+OrderedType = Union[None, bool, object]

pandas/core/arrays/categorical.py

@@ -1,5 +1,6 @@
 from shutil import get_terminal_size
 import textwrap
+from typing import Type, Union, cast
 from warnings import warn
 
 import numpy as np
@@ -47,6 +48,7 @@
 from pandas.core.dtypes.inference import is_hashable
 from pandas.core.dtypes.missing import isna, notna
 
+from pandas._typing import ArrayLike, Dtype, OrderedType
 from pandas.core import ops
 from pandas.core.accessor import PandasDelegate, delegate_names
 import pandas.core.algorithms as algorithms
@@ -466,7 +468,7 @@ def categories(self, categories):
         self._dtype = new_dtype
 
     @property
-    def ordered(self):
+    def ordered(self) -> OrderedType:
         """
         Whether the categories have an ordered relationship.
         """
@@ -480,11 +482,11 @@ def dtype(self) -> CategoricalDtype:
         return self._dtype
 
     @property
-    def _ndarray_values(self):
+    def _ndarray_values(self) -> np.ndarray:
         return self.codes
 
     @property
-    def _constructor(self):
+    def _constructor(self) -> Type["Categorical"]:
         return Categorical
 
     @classmethod
@@ -495,15 +497,15 @@ def _formatter(self, boxed=False):
         # Defer to CategoricalFormatter's formatter.
         return None
 
-    def copy(self):
+    def copy(self) -> "Categorical":
         """
         Copy constructor.
         """
         return self._constructor(
             values=self._codes.copy(), dtype=self.dtype, fastpath=True
         )
 
-    def astype(self, dtype, copy=True):
+    def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike:
         """
         Coerce this type to another dtype
 
@@ -519,6 +521,8 @@ def astype(self, dtype, copy=True):
 
         """
         if is_categorical_dtype(dtype):
+            dtype = cast(Union[str, CategoricalDtype], dtype)
+
             # GH 10696/18593
             dtype = self.dtype.update_dtype(dtype)
             self = self.copy() if copy else self
@@ -528,27 +532,27 @@ def astype(self, dtype, copy=True):
         return np.array(self, dtype=dtype, copy=copy)
 
     @cache_readonly
-    def ndim(self):
+    def ndim(self) -> int:
         """
         Number of dimensions of the Categorical
         """
         return self._codes.ndim
 
     @cache_readonly
-    def size(self):
+    def size(self) -> int:
         """
         return the len of myself
         """
         return len(self)
 
     @cache_readonly
-    def itemsize(self):
+    def itemsize(self) -> int:
         """
         return the size of a single category
         """
         return self.categories.itemsize
 
-    def tolist(self):
+    def tolist(self) -> list:
         """
         Return a list of the values.
 
@@ -561,7 +565,7 @@ def tolist(self):
     to_list = tolist
 
     @property
-    def base(self):
+    def base(self) -> None:
         """
         compat, we are always our own object
         """
@@ -769,7 +773,7 @@ def _set_categories(self, categories, fastpath=False):
 
         self._dtype = new_dtype
 
-    def _set_dtype(self, dtype):
+    def _set_dtype(self, dtype: CategoricalDtype) -> "Categorical":
         """
         Internal method for directly updating the CategoricalDtype
 

pandas/core/dtypes/dtypes.py

@@ -1,6 +1,6 @@
 """ define extension dtypes """
 import re
-from typing import Any, Dict, List, Optional, Tuple, Type, Union
+from typing import Any, Dict, List, Optional, Tuple, Type, Union, cast
 import warnings
 
 import numpy as np
@@ -11,6 +11,8 @@
 
 from pandas.core.dtypes.generic import ABCCategoricalIndex, ABCDateOffset, ABCIndexClass
 
+from pandas._typing import OrderedType
+
 from .base import ExtensionDtype
 from .inference import is_list_like
 
@@ -20,9 +22,6 @@
 # CategoricalDtype constructor to detect when ordered=None is explicitly passed
 ordered_sentinel = object()  # type: object
 
-# TODO(GH26403): Replace with Optional[bool] or bool
-OrderedType = Union[None, bool, object]
-
 
 def register_extension_dtype(cls: Type[ExtensionDtype],) -> Type[ExtensionDtype]:
     """
@@ -524,7 +523,9 @@ def validate_categories(categories, fastpath: bool = False):
 
         return categories
 
-    def update_dtype(self, dtype: "CategoricalDtype") -> "CategoricalDtype":
+    def update_dtype(
+        self, dtype: Union[str_type, "CategoricalDtype"]
+    ) -> "CategoricalDtype":
         """
         Returns a CategoricalDtype with categories and ordered taken from dtype
         if specified, otherwise falling back to self if unspecified
@@ -547,6 +548,8 @@ def update_dtype(self, dtype: "CategoricalDtype") -> "CategoricalDtype":
             ).format(dtype=dtype)
             raise ValueError(msg)
 
+        dtype = cast(CategoricalDtype, dtype)
+
         # dtype is CDT: keep current categories/ordered if None
         new_categories = dtype.categories
         if new_categories is None: