POC: 2D EAs via composition

pandas/core/algorithms.py

@@ -105,6 +105,9 @@ def _ensure_data(values, dtype=None):
         else:
             # Datetime
             from pandas import DatetimeIndex
+            from pandas.core.arrays import unwrap_reshapeable
+            values = unwrap_reshapeable(values)
+            assert values.ndim == 1, (type(values), values.shape)
             values = DatetimeIndex(values)
             dtype = values.dtype
 
@@ -1525,7 +1528,7 @@ def take(arr, indices, axis=0, allow_fill=False, fill_value=None):
 
     if allow_fill:
         # Pandas style, -1 means NA
-        validate_indices(indices, len(arr))
+        validate_indices(indices, arr.shape[axis])
         result = take_1d(arr, indices, axis=axis, allow_fill=True,
                          fill_value=fill_value)
     else:
@@ -1575,7 +1578,11 @@ def take_nd(arr, indexer, axis=0, out=None, fill_value=np.nan, mask_info=None,
     # TODO(EA): Remove these if / elifs as datetimeTZ, interval, become EAs
     # dispatch to internal type takes
     if is_extension_array_dtype(arr):
-        return arr.take(indexer, fill_value=fill_value, allow_fill=allow_fill)
+        try:
+            return arr.take(indexer, fill_value=fill_value, allow_fill=allow_fill, axis=axis)
+        except TypeError:
+            # `axis` kwarg not yet available
+            return arr.take(indexer, fill_value=fill_value, allow_fill=allow_fill)
     elif is_datetime64tz_dtype(arr):
         return arr.take(indexer, fill_value=fill_value, allow_fill=allow_fill)
     elif is_interval_dtype(arr):

pandas/core/arrays/__init__.py

@@ -11,3 +11,5 @@
     IntegerArray, integer_array)
 from .sparse import SparseArray  # noqa
 from .numpy_ import PandasArray, PandasDtype  # noqa
+from .reshaping import (  # noqa
+	ReshapeableArray, ReshapeMixin, unwrap_reshapeable)

pandas/core/arrays/base.py

@@ -111,10 +111,12 @@ class ExtensionArray:
     # '_typ' is for pandas.core.dtypes.generic.ABCExtensionArray.
     # Don't override this.
     _typ = 'extension'
+    _allows_2d = False
 
     # ------------------------------------------------------------------------
     # Constructors
     # ------------------------------------------------------------------------
+
     @classmethod
     def _from_sequence(cls, scalars, dtype=None, copy=False):
         """
@@ -286,6 +288,7 @@ def __iter__(self):
     # ------------------------------------------------------------------------
     # Required attributes
     # ------------------------------------------------------------------------
+
     @property
     def dtype(self) -> ExtensionDtype:
         """
@@ -305,7 +308,14 @@ 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.
+        """
+        return np.prod(self.shape)
 
     @property
     def nbytes(self) -> int:
@@ -319,6 +329,7 @@ def nbytes(self) -> int:
     # ------------------------------------------------------------------------
     # Additional Methods
     # ------------------------------------------------------------------------
+
     def astype(self, dtype, copy=True):
         """
         Cast to a NumPy array with 'dtype'.
@@ -479,8 +490,7 @@ def dropna(self):
     def shift(
             self,
             periods: int = 1,
-            fill_value: object = None,
-    ) -> ABCExtensionArray:
+            fill_value: object = None) -> ABCExtensionArray:
         """
         Shift values by desired number.
 
@@ -836,6 +846,7 @@ def copy(self, deep: bool = False) -> ABCExtensionArray:
     # ------------------------------------------------------------------------
     # Printing
     # ------------------------------------------------------------------------
+
     def __repr__(self):
         from pandas.io.formats.printing import format_object_summary
 

pandas/core/arrays/categorical.py

@@ -37,6 +37,8 @@
 from pandas.io.formats import console
 
 from .base import ExtensionArray, _extension_array_shared_docs
+from .reshaping import unwrap_reshapeable
+
 
 _take_msg = textwrap.dedent("""\
     Interpreting negative values in 'indexer' as missing values.
@@ -349,6 +351,7 @@ def __init__(self, values, categories=None, ordered=None, dtype=None,
                     values = [values[idx] for idx in np.where(~null_mask)[0]]
                 values = sanitize_array(values, None, dtype=sanitize_dtype)
 
+        values = unwrap_reshapeable(values)
         if dtype.categories is None:
             try:
                 codes, categories = factorize(values, sort=True)
@@ -457,11 +460,14 @@ def _formatter(self, boxed=False):
         # Defer to CategoricalFormatter's formatter.
         return None
 
-    def copy(self):
+    def copy(self, deep: bool = False):
         """
         Copy constructor.
         """
-        return self._constructor(values=self._codes.copy(),
+        values = self._codes
+        if deep:
+            values = values.copy()
+        return self._constructor(values=values,
                                  dtype=self.dtype,
                                  fastpath=True)
 

pandas/core/arrays/datetimelike.py

@@ -39,7 +39,7 @@
 from .base import ExtensionArray, ExtensionOpsMixin
 
 
-class AttributesMixin:
+class AttributesMixin:  # TODO: how much of this do we still need?
     _data = None  # type: np.ndarray
 
     @property
@@ -401,11 +401,6 @@ def __array__(self, dtype=None):
             return np.array(list(self), dtype=object)
         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)
 

pandas/core/arrays/datetimes.py

@@ -323,6 +323,9 @@ def __init__(self, values, dtype=_NS_DTYPE, freq=None, copy=False):
             )
             raise ValueError(msg.format(values.dtype))
 
+        if values.ndim != 1:
+            raise ValueError("Only 1-dimensional inputs are valid.")
+
         dtype = _validate_dt64_dtype(dtype)
 
         if freq == "infer":
@@ -353,7 +356,7 @@ def __init__(self, values, dtype=_NS_DTYPE, freq=None, copy=False):
 
     @classmethod
     def _simple_new(cls, values, freq=None, dtype=_NS_DTYPE):
-        assert isinstance(values, np.ndarray)
+        assert isinstance(values, np.ndarray), type(values)
         if values.dtype == 'i8':
             values = values.view(_NS_DTYPE)