TYPING: Partial typing of Categorical (#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.
  """