EA: support basic 2D operations

doc/source/whatsnew/v0.25.0.rst

@@ -567,6 +567,7 @@ Other API changes
 - Using an unsupported version of Beautiful Soup 4 will now raise an ``ImportError`` instead of a ``ValueError`` (:issue:`27063`)
 - :meth:`Series.to_excel` and :meth:`DataFrame.to_excel` will now raise a ``ValueError`` when saving timezone aware data. (:issue:`27008`, :issue:`7056`)
 - :meth:`DataFrame.to_hdf` and :meth:`Series.to_hdf` will now raise a ``NotImplementedError`` when saving a :class:`MultiIndex` with extention data types for a ``fixed`` format. (:issue:`7775`)
+- :meth:`Categorical.ravel` will now return a :class:`Categorical` instead of a NumPy array. (:issue:`27153`)
 
 .. _whatsnew_0250.deprecations:
 

pandas/core/arrays/__init__.py

@@ -1,3 +1,4 @@
+from ._reshaping import implement_2d  # noqa:F401
 from .array_ import array  # noqa: F401
 from .base import (  # noqa: F401
     ExtensionArray, ExtensionOpsMixin, ExtensionScalarOpsMixin)

pandas/core/arrays/_reshaping.py

@@ -0,0 +1,91 @@
+"""
+Utilities for implementing 2D compatibility for 1D ExtensionArrays.
+"""
+from typing import Tuple
+
+import numpy as np
+
+from pandas._libs.lib import is_integer
+
+
+def implement_2d(cls):
+    """
+    A decorator to take a 1-dimension-only ExtensionArray subclass and make
+    it support limited 2-dimensional operations.
+    """
+    from pandas.core.arrays import ExtensionArray
+
+    # For backwards-compatibility, if an EA author implemented __len__
+    #  but not size, we use that __len__ method to get an array's size.
+    has_size = cls.size is not ExtensionArray.size
+    has_shape = cls.shape is not ExtensionArray.shape
+    has_len = cls.__len__ is not ExtensionArray.__len__
+
+    if not has_size and has_len:
+        cls.size = property(cls.__len__)
+        cls.__len__ = ExtensionArray.__len__
+
+    elif not has_size and has_shape:
+        @property
+        def size(self) -> int:
+            return np.prod(self.shape)
+
+        cls.size = size
+
+    return cls
+
+
+def tuplify_shape(size: int, shape) -> Tuple[int, ...]:
+    """
+    Convert a passed shape into a valid tuple.
+    Following ndarray.reshape, we accept either `reshape(a, b)` or
+    `reshape((a, b))`, the latter being canonical.
+
+    Parameters
+    ----------
+    size : int
+    shape : tuple
+
+    Returns
+    -------
+    tuple[int, ...]
+    """
+    if len(shape) == 0:
+        raise ValueError("shape must be a non-empty tuple of integers",
+                         shape)
+
+    if len(shape) == 1:
+        if is_integer(shape[0]):
+            pass
+        else:
+            shape = shape[0]
+            if not isinstance(shape, tuple):
+                raise ValueError("shape must be a non-empty tuple of integers",
+                                 shape)
+
+    if not all(is_integer(x) for x in shape):
+        raise ValueError("shape must be a non-empty tuple of integers", shape)
+
+    if any(x < -1 for x in shape):
+        raise ValueError("Invalid shape {shape}".format(shape=shape))
+
+    if -1 in shape:
+        if shape.count(-1) != 1:
+            raise ValueError("Invalid shape {shape}".format(shape=shape))
+        idx = shape.index(-1)
+        others = [n for n in shape if n != -1]
+        prod = np.prod(others)
+        dim = size // prod
+        shape = shape[:idx] + (dim,) + shape[idx + 1:]
+
+    if np.prod(shape) != size:
+        raise ValueError("Product of shape ({shape}) must match "
+                         "size ({size})".format(shape=shape,
+                                                size=size))
+
+    num_gt1 = len([x for x in shape if x > 1])
+    if num_gt1 > 1:
+        raise ValueError("The default `reshape` implementation is limited to "
+                         "shapes (N,), (N,1), and (1,N), not {shape}"
+                         .format(shape=shape))
+    return shape

pandas/core/arrays/base.py

@@ -23,6 +23,7 @@
 
 from pandas._typing import ArrayLike
 from pandas.core import ops
+from pandas.core.arrays._reshaping import tuplify_shape
 
 _not_implemented_message = "{} does not implement {}."
 
@@ -47,12 +48,13 @@ class ExtensionArray:
     * _from_sequence
     * _from_factorized
     * __getitem__
-    * __len__
+    * __len__ *or* size
     * dtype
     * nbytes
     * isna
     * take
     * copy
+    * view
     * _concat_same_type
 
     A default repr displaying the type, (truncated) data, length,
@@ -123,6 +125,12 @@ class ExtensionArray:
     # Don't override this.
     _typ = 'extension'
 
+    # Whether this class supports 2D arrays natively.  If so, set _allows_2d
+    #  to True and override reshape, ravel, and T.  Otherwise, apply the
+    #  `implement_2d` decorator to use default implementations of limited
+    #  2D functionality.
+    _allows_2d = False
+
     # ------------------------------------------------------------------------
     # Constructors
     # ------------------------------------------------------------------------
@@ -283,7 +291,7 @@ def __len__(self) -> int:
         -------
         length : int
         """
-        raise AbstractMethodError(self)
+        return self.shape[0]
 
     def __iter__(self):
         """
@@ -298,6 +306,7 @@ def __iter__(self):
     # ------------------------------------------------------------------------
     # Required attributes
     # ------------------------------------------------------------------------
+    _shape = None
 
     @property
     def dtype(self) -> ExtensionDtype:
@@ -311,14 +320,26 @@ def shape(self) -> Tuple[int, ...]:
         """
         Return a tuple of the array dimensions.
         """
-        return (len(self),)
+        if self._shape is not None:
+            return self._shape
+
+        # Default to 1D
+        length = self.size
+        return (length,)
 
     @property
     def ndim(self) -> int:
         """
         Extension Arrays are only allowed to be 1-dimensional.
         """
-        return 1
+        return len(self.shape)
+
+    @property
+    def size(self) -> int:
+        """
+        The number of elements in this array.
+        """
+        raise AbstractMethodError(self)
 
     @property
     def nbytes(self) -> int:
@@ -841,6 +862,22 @@ def copy(self) -> ABCExtensionArray:
         """
         raise AbstractMethodError(self)
 
+    def view(self, dtype=None) -> ABCExtensionArray:
+        """
+        Return a view on the array.
+
+        Returns
+        -------
+        ExtensionArray
+
+        Notes
+        -----
+        - This must return a *new* object, not self.
+        - The only case that *must* be implemented is with dtype=None,
+          giving a view with the same dtype as self.
+        """
+        raise AbstractMethodError(self)
+
     # ------------------------------------------------------------------------
     # Printing
     # ------------------------------------------------------------------------
@@ -908,6 +945,30 @@ def _formatting_values(self) -> np.ndarray:
     # Reshaping
     # ------------------------------------------------------------------------
 
+    def reshape(self, *shape):
+        # numpy accepts either a single tuple or an expanded tuple
+        shape = tuplify_shape(self.size, shape)
+        result = self.view()
+        result._shape = shape
+        return result
+
+    @property
+    def T(self) -> ABCExtensionArray:
+        """
+        Return a transposed view on self.
+        """
+        shape = self.shape[::-1]
+        return self.reshape(shape)
+
+    def ravel(self, order=None) -> ABCExtensionArray:
+        """
+        Return a flattened view on self.
+        """
+        # Note: we ignore `order`, keep the argument for compat with
+        #  numpy signature.
+        shape = (self.size,)
+        return self.reshape(shape)
+
     @classmethod
     def _concat_same_type(
             cls,

pandas/core/arrays/categorical.py

@@ -37,6 +37,7 @@
 
 from pandas.io.formats import console
 
+from ._reshaping import implement_2d
 from .base import ExtensionArray, _extension_array_shared_docs
 
 _take_msg = textwrap.dedent("""\
@@ -208,6 +209,7 @@ def contains(cat, key, container):
 """
 
 
+@implement_2d
 class Categorical(ExtensionArray, PandasObject):
     """
     Represent a categorical variable in classic R / S-plus fashion.
@@ -490,20 +492,6 @@ def astype(self, dtype, copy=True):
             return self._set_dtype(dtype)
         return np.array(self, dtype=dtype, copy=copy)
 
-    @cache_readonly
-    def ndim(self):
-        """
-        Number of dimensions of the Categorical
-        """
-        return self._codes.ndim
-
-    @cache_readonly
-    def size(self):
-        """
-        return the len of myself
-        """
-        return len(self)
-
     @cache_readonly
     def itemsize(self):
         """
@@ -1230,8 +1218,7 @@ def shape(self):
         -------
         shape : tuple
         """
-
-        return tuple([len(self._codes)])
+        return self._codes.shape
 
     def shift(self, periods, fill_value=None):
         """
@@ -1697,19 +1684,7 @@ def _values_for_rank(self):
             )
         return values
 
-    def ravel(self, order='C'):
-        """
-        Return a flattened (numpy) array.
-
-        For internal compatibility with numpy arrays.
-
-        Returns
-        -------
-        numpy.array
-        """
-        return np.array(self)
-
-    def view(self):
+    def view(self, dtype=None):
         """
         Return a view of myself.
 
@@ -1720,7 +1695,11 @@ def view(self):
         view : Categorical
            Returns `self`!
         """
-        return self
+        if dtype is not None:
+            return NotImplementedError(dtype)
+        return self._constructor(values=self._codes,
+                                 dtype=self.dtype,
+                                 fastpath=True)
 
     def to_dense(self):
         """
@@ -1935,12 +1914,6 @@ def _slice(self, slicer):
         codes = self._codes[slicer]
         return self._constructor(values=codes, dtype=self.dtype, fastpath=True)
 
-    def __len__(self):
-        """
-        The length of this Categorical.
-        """
-        return len(self._codes)
-
     def __iter__(self):
         """
         Returns an Iterator over the values of this Categorical.

pandas/core/arrays/datetimelike.py

@@ -36,6 +36,7 @@
 from pandas.tseries import frequencies
 from pandas.tseries.offsets import DateOffset, Tick
 
+from ._reshaping import implement_2d
 from .base import ExtensionArray, ExtensionOpsMixin
 
 
@@ -324,6 +325,7 @@ def ceil(self, freq, ambiguous='raise', nonexistent='raise'):
         return self._round(freq, RoundTo.PLUS_INFTY, ambiguous, nonexistent)
 
 
+@implement_2d
 class DatetimeLikeArrayMixin(ExtensionOpsMixin,
                              AttributesMixin,
                              ExtensionArray):
@@ -402,12 +404,8 @@ def __array__(self, dtype=None):
         return self._data
 
     @property
-    def size(self) -> int:
-        """The number of elements in this array."""
-        return np.prod(self.shape)
-
-    def __len__(self):
-        return len(self._data)
+    def shape(self):
+        return self._data.shape
 
     def __getitem__(self, key):
         """
@@ -557,6 +555,8 @@ def view(self, dtype=None):
         ndarray
             With the specified `dtype`.
         """
+        if dtype is None:
+            return type(self)(self._data, dtype=self.dtype, freq=self.freq)
         return self._data.view(dtype=dtype)
 
     # ------------------------------------------------------------------