TYPING: Partial typing of Categorical (pandas-dev#27318)

Author: topper-123

pandas/_typing.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import IO, TYPE_CHECKING, AnyStr, TypeVar, Union
+from typing import IO, TYPE_CHECKING, AnyStr, Optional, TypeVar, Union
 
 import numpy as np
 
@@ -27,3 +27,4 @@
 FrameOrSeries = TypeVar("FrameOrSeries", "Series", "DataFrame")
 Scalar = Union[str, int, float]
 Axis = Union[str, int]
+Ordered = Optional[bool]

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
@@ -42,6 +43,7 @@
 from pandas.core.dtypes.inference import is_hashable
 from pandas.core.dtypes.missing import isna, notna
 
+from pandas._typing import ArrayLike, Dtype, Ordered
 from pandas.core import ops
 from pandas.core.accessor import PandasDelegate, delegate_names
 import pandas.core.algorithms as algorithms
@@ -455,7 +457,7 @@ def categories(self, categories):
         self._dtype = new_dtype
 
     @property
-    def ordered(self):
+    def ordered(self) -> Ordered:
         """
         Whether the categories have an ordered relationship.
         """
@@ -469,11 +471,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
@@ -484,15 +486,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
 
@@ -505,6 +507,8 @@ def astype(self, dtype, copy=True):
             object is returned.
         """
         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
@@ -514,27 +518,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.
 
@@ -547,7 +551,7 @@ def tolist(self):
     to_list = tolist
 
     @property
-    def base(self):
+    def base(self) -> None:
         """
         compat, we are always our own object
         """
@@ -694,8 +698,6 @@ def from_codes(cls, codes, categories=None, ordered=None, dtype=None):
 
         return cls(codes, dtype=dtype, fastpath=True)
 
-    _codes = None
-
     def _get_codes(self):
         """
         Get the codes.
@@ -755,7 +757,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 Ordered
+
 from .base import ExtensionDtype
 from .inference import is_bool, 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]:
     """
@@ -222,7 +221,11 @@ class CategoricalDtype(PandasExtensionDtype, ExtensionDtype):
     _metadata = ("categories", "ordered", "_ordered_from_sentinel")
     _cache = {}  # type: Dict[str_type, PandasExtensionDtype]
 
-    def __init__(self, categories=None, ordered: OrderedType = ordered_sentinel):
+    def __init__(
+        self, categories=None, ordered: Union[Ordered, object] = ordered_sentinel
+    ):
+        # TODO(GH26403): Set type of ordered to Ordered
+        ordered = cast(Ordered, ordered)
         self._finalize(categories, ordered, fastpath=False)
 
     @classmethod
@@ -235,7 +238,7 @@ def _from_fastpath(
 
     @classmethod
     def _from_categorical_dtype(
-        cls, dtype: "CategoricalDtype", categories=None, ordered: OrderedType = None
+        cls, dtype: "CategoricalDtype", categories=None, ordered: Ordered = None
     ) -> "CategoricalDtype":
         if categories is ordered is None:
             return dtype
@@ -336,9 +339,7 @@ def _from_values_or_dtype(
 
         return dtype
 
-    def _finalize(
-        self, categories, ordered: OrderedType, fastpath: bool = False
-    ) -> None:
+    def _finalize(self, categories, ordered: Ordered, fastpath: bool = False) -> None:
 
         if ordered is not None and ordered is not ordered_sentinel:
             self.validate_ordered(ordered)
@@ -423,7 +424,7 @@ def __repr__(self):
         return tpl.format(data, self._ordered)
 
     @staticmethod
-    def _hash_categories(categories, ordered: OrderedType = True) -> int:
+    def _hash_categories(categories, ordered: Ordered = True) -> int:
         from pandas.core.util.hashing import (
             hash_array,
             _combine_hash_arrays,
@@ -475,7 +476,7 @@ def construct_array_type(cls):
         return Categorical
 
     @staticmethod
-    def validate_ordered(ordered: OrderedType) -> None:
+    def validate_ordered(ordered: Ordered) -> None:
         """
         Validates that we have a valid ordered parameter. If
         it is not a boolean, a TypeError will be raised.
@@ -529,7 +530,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
@@ -551,6 +554,9 @@ def update_dtype(self, dtype: "CategoricalDtype") -> "CategoricalDtype":
                 "got {dtype!r}"
             ).format(dtype=dtype)
             raise ValueError(msg)
+        else:
+            # from here on, dtype is a CategoricalDtype
+            dtype = cast(CategoricalDtype, dtype)
 
         # dtype is CDT: keep current categories/ordered if None
         new_categories = dtype.categories
@@ -583,7 +589,7 @@ def categories(self):
         return self._categories
 
     @property
-    def ordered(self) -> OrderedType:
+    def ordered(self) -> Ordered:
         """
         Whether the categories have an ordered relationship.
         """