REF: mix NDArrayBacked into Categorical, PandasArray (#41131)

Author: jbrockmendel

pandas/_libs/arrays.pyi

@@ -0,0 +1,45 @@
+from typing import Sequence
+
+import numpy as np
+
+from pandas._typing import (
+    DtypeObj,
+    Shape,
+)
+
+class NDArrayBacked:
+    _dtype: DtypeObj
+    _ndarray: np.ndarray
+
+    def __init__(self, values: np.ndarray, dtype: DtypeObj): ...
+
+    @classmethod
+    def _simple_new(cls, values: np.ndarray, dtype: DtypeObj): ...
+
+    def _from_backing_data(self, values: np.ndarray): ...
+
+    def __setstate__(self, state): ...
+
+    def __len__(self) -> int: ...
+
+    @property
+    def shape(self) -> Shape: ...
+
+    @property
+    def ndim(self) -> int: ...
+
+    @property
+    def size(self) -> int: ...
+
+    @property
+    def nbytes(self) -> int: ...
+
+    def copy(self): ...
+    def delete(self, loc, axis=0): ...
+    def swapaxes(self, axis1, axis2): ...
+    def repeat(self, repeats: int | Sequence[int], axis: int | None = ...): ...
+    def reshape(self, *args, **kwargs): ...
+    def ravel(self, order="C"): ...
+
+    @property
+    def T(self): ...

pandas/_libs/tslibs/dtypes.pyi

@@ -7,6 +7,8 @@ _period_code_map: dict[str, int]
 
 
 class PeriodDtypeBase:
+    _dtype_code: int  # PeriodDtypeCode
+
     # actually __cinit__
     def __new__(self, code: int): ...
 

pandas/core/arrays/_mixins.py

@@ -11,18 +11,15 @@
 import numpy as np
 
 from pandas._libs import lib
+from pandas._libs.arrays import NDArrayBacked
 from pandas._typing import (
     F,
     PositionalIndexer2D,
     Shape,
     type_t,
 )
-from pandas.compat.numpy import function as nv
 from pandas.errors import AbstractMethodError
-from pandas.util._decorators import (
-    cache_readonly,
-    doc,
-)
+from pandas.util._decorators import doc
 from pandas.util._validators import (
     validate_bool_kwarg,
     validate_fillna_kwargs,
@@ -69,24 +66,13 @@ def method(self, *args, **kwargs):
     return cast(F, method)
 
 
-class NDArrayBackedExtensionArray(ExtensionArray):
+class NDArrayBackedExtensionArray(NDArrayBacked, ExtensionArray):
     """
     ExtensionArray that is backed by a single NumPy ndarray.
     """
 
     _ndarray: np.ndarray
 
-    def _from_backing_data(
-        self: NDArrayBackedExtensionArrayT, arr: np.ndarray
-    ) -> NDArrayBackedExtensionArrayT:
-        """
-        Construct a new ExtensionArray `new_array` with `arr` as its _ndarray.
-
-        This should round-trip:
-            self == self._from_backing_data(self._ndarray)
-        """
-        raise AbstractMethodError(self)
-
     def _box_func(self, x):
         """
         Wrap numpy type in our dtype.type if necessary.
@@ -142,46 +128,6 @@ def _validate_fill_value(self, fill_value):
 
     # ------------------------------------------------------------------------
 
-    # TODO: make this a cache_readonly; for that to work we need to remove
-    #  the _index_data kludge in libreduction
-    @property
-    def shape(self) -> Shape:
-        return self._ndarray.shape
-
-    def __len__(self) -> int:
-        return self.shape[0]
-
-    @cache_readonly
-    def ndim(self) -> int:
-        return len(self.shape)
-
-    @cache_readonly
-    def size(self) -> int:
-        return self._ndarray.size
-
-    @cache_readonly
-    def nbytes(self) -> int:
-        return self._ndarray.nbytes
-
-    def reshape(
-        self: NDArrayBackedExtensionArrayT, *args, **kwargs
-    ) -> NDArrayBackedExtensionArrayT:
-        new_data = self._ndarray.reshape(*args, **kwargs)
-        return self._from_backing_data(new_data)
-
-    def ravel(
-        self: NDArrayBackedExtensionArrayT, *args, **kwargs
-    ) -> NDArrayBackedExtensionArrayT:
-        new_data = self._ndarray.ravel(*args, **kwargs)
-        return self._from_backing_data(new_data)
-
-    @property
-    def T(self: NDArrayBackedExtensionArrayT) -> NDArrayBackedExtensionArrayT:
-        new_data = self._ndarray.T
-        return self._from_backing_data(new_data)
-
-    # ------------------------------------------------------------------------
-
     def equals(self, other) -> bool:
         if type(self) is not type(other):
             return False
@@ -208,24 +154,6 @@ def argmax(self, axis: int = 0, skipna: bool = True):  # type:ignore[override]
             raise NotImplementedError
         return nargminmax(self, "argmax", axis=axis)
 
-    def copy(self: NDArrayBackedExtensionArrayT) -> NDArrayBackedExtensionArrayT:
-        new_data = self._ndarray.copy()
-        return self._from_backing_data(new_data)
-
-    def repeat(
-        self: NDArrayBackedExtensionArrayT, repeats, axis=None
-    ) -> NDArrayBackedExtensionArrayT:
-        """
-        Repeat elements of an array.
-
-        See Also
-        --------
-        numpy.ndarray.repeat
-        """
-        nv.validate_repeat((), {"axis": axis})
-        new_data = self._ndarray.repeat(repeats, axis=axis)
-        return self._from_backing_data(new_data)
-
     def unique(self: NDArrayBackedExtensionArrayT) -> NDArrayBackedExtensionArrayT:
         new_data = unique(self._ndarray)
         return self._from_backing_data(new_data)
@@ -418,18 +346,6 @@ def where(
         res_values = np.where(mask, self._ndarray, value)
         return self._from_backing_data(res_values)
 
-    def delete(
-        self: NDArrayBackedExtensionArrayT, loc, axis: int = 0
-    ) -> NDArrayBackedExtensionArrayT:
-        res_values = np.delete(self._ndarray, loc, axis=axis)
-        return self._from_backing_data(res_values)
-
-    def swapaxes(
-        self: NDArrayBackedExtensionArrayT, axis1, axis2
-    ) -> NDArrayBackedExtensionArrayT:
-        res_values = self._ndarray.swapaxes(axis1, axis2)
-        return self._from_backing_data(res_values)
-
     # ------------------------------------------------------------------------
     # Additional array methods
     #  These are not part of the EA API, but we implement them because

pandas/core/arrays/categorical.py

@@ -27,6 +27,7 @@
     algos as libalgos,
     hashtable as htable,
 )
+from pandas._libs.arrays import NDArrayBacked
 from pandas._libs.lib import no_default
 from pandas._typing import (
     ArrayLike,
@@ -349,12 +350,13 @@ class Categorical(NDArrayBackedExtensionArray, PandasObject, ObjectStringArrayMi
     # For comparisons, so that numpy uses our implementation if the compare
     # ops, which raise
     __array_priority__ = 1000
-    _dtype = CategoricalDtype(ordered=False)
     # tolist is not actually deprecated, just suppressed in the __dir__
     _hidden_attrs = PandasObject._hidden_attrs | frozenset(["tolist"])
     _typ = "categorical"
     _can_hold_na = True
 
+    _dtype: CategoricalDtype
+
     def __init__(
         self,
         values,
@@ -373,8 +375,9 @@ def __init__(
         # infer categories in a factorization step further below
 
         if fastpath:
-            self._ndarray = coerce_indexer_dtype(values, dtype.categories)
-            self._dtype = self._dtype.update_dtype(dtype)
+            codes = coerce_indexer_dtype(values, dtype.categories)
+            dtype = CategoricalDtype(ordered=False).update_dtype(dtype)
+            super().__init__(codes, dtype)
             return
 
         if not is_list_like(values):
@@ -463,8 +466,11 @@ def __init__(
             full_codes[~null_mask] = codes
             codes = full_codes
 
-        self._dtype = self._dtype.update_dtype(dtype)
-        self._ndarray = coerce_indexer_dtype(codes, dtype.categories)
+        dtype = CategoricalDtype(ordered=False).update_dtype(dtype)
+        arr = coerce_indexer_dtype(codes, dtype.categories)
+        # error: Argument 1 to "__init__" of "NDArrayBacked" has incompatible
+        # type "Union[ExtensionArray, ndarray]"; expected "ndarray"
+        super().__init__(arr, dtype)  # type: ignore[arg-type]
 
     @property
     def dtype(self) -> CategoricalDtype:
@@ -513,9 +519,7 @@ def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike:
             raise ValueError("Cannot convert float NaN to integer")
 
         elif len(self.codes) == 0 or len(self.categories) == 0:
-            # error: Incompatible types in assignment (expression has type "ndarray",
-            # variable has type "Categorical")
-            result = np.array(  # type: ignore[assignment]
+            result = np.array(
                 self,
                 dtype=dtype,
                 copy=copy,
@@ -533,11 +537,7 @@ def astype(self, dtype: Dtype, copy: bool = True) -> ArrayLike:
                 msg = f"Cannot cast {self.categories.dtype} dtype to {dtype}"
                 raise ValueError(msg)
 
-            # error: Incompatible types in assignment (expression has type "ndarray",
-            # variable has type "Categorical")
-            result = take_nd(  # type: ignore[assignment]
-                new_cats, ensure_platform_int(self._codes)
-            )
+            result = take_nd(new_cats, ensure_platform_int(self._codes))
 
         return result
 
@@ -745,7 +745,7 @@ def categories(self, categories):
                 "new categories need to have the same number of "
                 "items as the old categories!"
             )
-        self._dtype = new_dtype
+        super().__init__(self._ndarray, new_dtype)
 
     @property
     def ordered(self) -> Ordered:
@@ -809,7 +809,7 @@ def _set_categories(self, categories, fastpath=False):
                 "items than the old categories!"
             )
 
-        self._dtype = new_dtype
+        super().__init__(self._ndarray, new_dtype)
 
     def _set_dtype(self, dtype: CategoricalDtype) -> Categorical:
         """
@@ -842,7 +842,7 @@ def set_ordered(self, value, inplace=False):
         inplace = validate_bool_kwarg(inplace, "inplace")
         new_dtype = CategoricalDtype(self.categories, ordered=value)
         cat = self if inplace else self.copy()
-        cat._dtype = new_dtype
+        NDArrayBacked.__init__(cat, cat._ndarray, new_dtype)
         if not inplace:
             return cat
 
@@ -961,12 +961,12 @@ def set_categories(
             ):
                 # remove all _codes which are larger and set to -1/NaN
                 cat._codes[cat._codes >= len(new_dtype.categories)] = -1
+            codes = cat._codes
         else:
             codes = recode_for_categories(
                 cat.codes, cat.categories, new_dtype.categories
             )
-            cat._ndarray = codes
-        cat._dtype = new_dtype
+        NDArrayBacked.__init__(cat, codes, new_dtype)
 
         if not inplace:
             return cat
@@ -1182,8 +1182,8 @@ def add_categories(self, new_categories, inplace=no_default):
         new_dtype = CategoricalDtype(new_categories, self.ordered)
 
         cat = self if inplace else self.copy()
-        cat._dtype = new_dtype
-        cat._ndarray = coerce_indexer_dtype(cat._ndarray, new_dtype.categories)
+        codes = coerce_indexer_dtype(cat._ndarray, new_dtype.categories)
+        NDArrayBacked.__init__(cat, codes, new_dtype)
         if not inplace:
             return cat
 
@@ -1303,9 +1303,8 @@ def remove_unused_categories(self, inplace=no_default):
         new_dtype = CategoricalDtype._from_fastpath(
             new_categories, ordered=self.ordered
         )
-        cat._dtype = new_dtype
-        cat._ndarray = coerce_indexer_dtype(inv, new_dtype.categories)
-
+        new_codes = coerce_indexer_dtype(inv, new_dtype.categories)
+        NDArrayBacked.__init__(cat, new_codes, new_dtype)
         if not inplace:
             return cat
 
@@ -1484,7 +1483,7 @@ def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
     def __setstate__(self, state):
         """Necessary for making this object picklable"""
         if not isinstance(state, dict):
-            raise Exception("invalid pickle state")
+            return super().__setstate__(state)
 
         if "_dtype" not in state:
             state["_dtype"] = CategoricalDtype(state["_categories"], state["_ordered"])
@@ -1493,8 +1492,7 @@ def __setstate__(self, state):
             # backward compat, changed what is property vs attribute
             state["_ndarray"] = state.pop("_codes")
 
-        for k, v in state.items():
-            setattr(self, k, v)
+        super().__setstate__(state)
 
     @property
     def nbytes(self) -> int:
@@ -1863,16 +1861,7 @@ def _codes(self) -> np.ndarray:
 
     @_codes.setter
     def _codes(self, value: np.ndarray):
-        self._ndarray = value
-
-    def _from_backing_data(self, arr: np.ndarray) -> Categorical:
-        assert isinstance(arr, np.ndarray)
-        assert arr.dtype == self._ndarray.dtype
-
-        res = object.__new__(type(self))
-        res._ndarray = arr
-        res._dtype = self.dtype
-        return res
+        NDArrayBacked.__init__(self, value, self.dtype)
 
     def _box_func(self, i: int):
         if i == -1: