Internals refactor

Author: TomAugspurger

pandas/core/arrays/__init__.py

@@ -1 +1,2 @@
-from .categorical import Categorical
+from .base import ExtensionArray  # noqa
+from .categorical import Categorical  # noqa

pandas/core/arrays/base.py

@@ -1,6 +1,6 @@
 """An interface for extending pandas with custom arrays."""
 import abc
-from typing import Tuple, Sequence, Optional, Any  # noqa
+from typing import Tuple, Sequence, Optional, Any, Iterator  # noqa
 
 import numpy as np
 
@@ -20,19 +20,23 @@ class ExtensionArray(metaclass=abc.ABCMeta):
     # ------------------------------------------------------------------------
     @abc.abstractmethod
     def __getitem__(self, item):
+        # type (Any) -> Any
         pass
 
     def __setitem__(self, key, value):
+        # type: (Any, Any) -> None
         raise NotImplementedError(_not_implemented_message.format(
             type(self), '__setitem__')
         )
 
     @abc.abstractmethod
     def __iter__(self):
+        # type: () -> Iterator
         pass
 
     @abc.abstractmethod
     def __len__(self):
+        # type: () -> int
         pass
 
     # ------------------------------------------------------------------------
@@ -41,6 +45,7 @@ def __len__(self):
     @property
     @abc.abstractmethod
     def dtype(self):
+        """An instance of 'ExtensionDtype'."""
         # type: () -> ExtensionDtype
         pass
 
@@ -52,23 +57,23 @@ def shape(self):
     @property
     def ndim(self):
         # type: () -> int
-        """Extension Arrays are only allowed to be 1-dimensional"""
+        """Extension Arrays are only allowed to be 1-dimensional."""
         return 1
 
     @property
     @abc.abstractmethod
     def nbytes(self):
+        """The number of bytes needed to store this object in memory."""
         # type: () -> int
-        # TODO: default impl?
         pass
 
     # ------------------------------------------------------------------------
     # Additional Methods
     # ------------------------------------------------------------------------
     @abc.abstractmethod
     def isna(self):
-        # type: () -> Sequence[bool]
-        # TODO: narrow this type?
+        """Boolean NumPy array indicating if each value is missing."""
+        # type: () -> np.ndarray
         pass
 
     # ------------------------------------------------------------------------

pandas/core/categorical.py

@@ -0,0 +1,2 @@
+# pickle compatability, probably a better way
+from pandas.core.arrays import Categorical  # noqa

pandas/core/internals.py

@@ -59,7 +59,7 @@
 
 from pandas.core.index import Index, MultiIndex, _ensure_index
 from pandas.core.indexing import maybe_convert_indices, length_of_indexer
-from pandas.core.arrays.categorical import Categorical, _maybe_to_categorical
+from pandas.core.arrays import Categorical, ExtensionArray
 from pandas.core.indexes.datetimes import DatetimeIndex
 from pandas.io.formats.printing import pprint_thing
 
@@ -95,6 +95,7 @@ class Block(PandasObject):
     is_object = False
     is_categorical = False
     is_sparse = False
+    is_extension = False
     _box_to_block_values = True
     _can_hold_na = False
     _downcast_dtype = None
@@ -108,14 +109,14 @@ class Block(PandasObject):
     def __init__(self, values, placement, ndim=None, fastpath=False):
         if ndim is None:
             ndim = values.ndim
-        elif values.ndim != ndim:
+        elif self._validate_ndim and values.ndim != ndim:
             raise ValueError('Wrong number of dimensions')
         self.ndim = ndim
 
         self.mgr_locs = placement
         self.values = values
 
-        if ndim and len(self.mgr_locs) != len(self.values):
+        if self._validate_ndim and ndim and len(self.mgr_locs) != len(self.values):
             raise ValueError(
                 'Wrong number of items passed {val}, placement implies '
                 '{mgr}'.format(val=len(self.values), mgr=len(self.mgr_locs)))
@@ -274,7 +275,6 @@ def reshape_nd(self, labels, shape, ref_items, mgr=None):
 
         return a new block that is transformed to a nd block
         """
-
         return _block2d_to_blocknd(values=self.get_values().T,
                                    placement=self.mgr_locs, shape=shape,
                                    labels=labels, ref_items=ref_items)
@@ -1697,24 +1697,19 @@ class NonConsolidatableMixIn(object):
     _holder = None
 
     def __init__(self, values, placement, ndim=None, fastpath=False, **kwargs):
+        # Placement must be converted to BlockPlacement so that we can check
+        # its length
+        if not isinstance(placement, BlockPlacement):
+            placement = BlockPlacement(placement)
 
-        # Placement must be converted to BlockPlacement via property setter
-        # before ndim logic, because placement may be a slice which doesn't
-        # have a length.
-        self.mgr_locs = placement
-
-        # kludgetastic
+        # Maybe infer ndim from placement
         if ndim is None:
-            if len(self.mgr_locs) != 1:
+            if len(placement) != 1:
                 ndim = 1
             else:
                 ndim = 2
-        self.ndim = ndim
-
-        if not isinstance(values, self._holder):
-            raise TypeError("values must be {0}".format(self._holder.__name__))
-
-        self.values = values
+        super(NonConsolidatableMixIn, self).__init__(values, placement, ndim=ndim,
+                                                     fastpath=fastpath)
 
     @property
     def shape(self):
@@ -1765,7 +1760,7 @@ def putmask(self, mask, new, align=True, inplace=False, axis=0,
 
         Returns
         -------
-        a new block(s), the result of the putmask
+        a new block, the result of the putmask
         """
         inplace = validate_bool_kwarg(inplace, 'inplace')
 
@@ -1823,6 +1818,91 @@ def _unstack(self, unstacker_func, new_columns):
         return blocks, mask
 
 
+class ExtensionBlock(NonConsolidatableMixIn, Block):
+    """Block for holding extension types.
+
+    Notes
+    -----
+    This is the holds all 3rd-party extension types. It's also the immediate
+    parent class for our internal extension types' blocks, CategoricalBlock.
+
+    All extension arrays *must* be 1-D, which simplifies things a bit.
+    """
+    # Some questions / notes as comments, will be removed.
+    #
+    # Currently inherited from NCB. We'll keep it around until SparseBlock
+    # and DatetimeTZBlock are refactored.
+    # - set
+    # - iget
+    # - should_store
+    # - putmask
+    # - _slice
+    # - _try_cast_result
+    # - unstack
+
+    # Think about overriding these methods from Block
+    # - _maybe_downcast: (never downcast)
+
+    # Methods we can (probably) ignore and just use Block's:
+
+    # * replace / replace_single
+    #   Categorical got Object, but was hopefully unnescessary.
+    #   DatetimeTZ, Sparse got Block
+    # * is_view
+    #   Categorical overrides to say that it is not.
+    #   DatetimeTZ, Sparse inherits Base anyway
+
+    is_extension = True
+
+    # XXX
+    # is_bool is is a change for CategoricalBlock. Used to inherit
+    # from Object to infer from values. If this matters, we should
+    # override it directly in CategoricalBlock so that we infer from
+    # the categories, not the codes.
+    is_bool = False
+
+    def __init__(self, values, placement, ndim=None, fastpath=False):
+        self._holder = type(values)
+        super(ExtensionBlock, self).__init__(values, placement, ndim=ndim,
+                                             fastpath=fastpath)
+
+    def get_values(self, dtype=None):
+        # ExtensionArrays must be iterable, so this works.
+        values = np.asarray(self.values)
+        if values.ndim == self.ndim - 1:
+            values = values.reshape((1,) + values.shape)
+        return values
+
+    def _can_hold_element(self, element):
+        # XXX:
+        # Not defined on NCM.
+        # Categorical got True from ObjectBlock
+        # DatetimeTZ gets DatetimeBlock
+        # Sparse gets Block
+        # Let's just assume yes for now, but we can maybe push
+        # this onto the array.
+        return True
+
+    def convert(self, copy=True, **kwargs):
+        # We're dedicated to a type, we don't convert.
+        # Taken from CategoricalBlock / Block.
+        return self.copy() if copy else self
+
+    def _slice(self, slicer):
+        """ return a slice of my values """
+
+        # slice the category
+        # return same dims as we currently have
+
+        if isinstance(slicer, tuple) and len(slicer) == 2:
+            if not is_null_slice(slicer[0]):
+                raise AssertionError("invalid slicing for a 1-ndim "
+                                     "categorical")
+            slicer = slicer[1]
+
+        return self.values._slice(slicer)
+
+
 class NumericBlock(Block):
     __slots__ = ()
     is_numeric = True
@@ -2337,7 +2417,7 @@ def re_replacer(s):
         return block
 
 
-class CategoricalBlock(NonConsolidatableMixIn, ObjectBlock):
+class CategoricalBlock(ExtensionBlock):
     __slots__ = ()
     is_categorical = True
     _verify_integrity = True
@@ -2346,6 +2426,7 @@ class CategoricalBlock(NonConsolidatableMixIn, ObjectBlock):
     _concatenator = staticmethod(_concat._concat_categorical)
 
     def __init__(self, values, placement, fastpath=False, **kwargs):
+        from pandas.core.arrays.categorical import _maybe_to_categorical
 
         # coerce to categorical if we can
         super(CategoricalBlock, self).__init__(_maybe_to_categorical(values),
@@ -2360,23 +2441,13 @@ def is_view(self):
     def to_dense(self):
         return self.values.to_dense().view()
 
-    def convert(self, copy=True, **kwargs):
-        return self.copy() if copy else self
-
     @property
     def array_dtype(self):
         """ the dtype to return if I want to construct this block as an
         array
         """
         return np.object_
 
-    def _slice(self, slicer):
-        """ return a slice of my values """
-
-        # slice the category
-        # return same dims as we currently have
-        return self.values._slice(slicer)
-
     def _try_coerce_result(self, result):
         """ reverse of try_coerce_args """
 
@@ -2468,7 +2539,8 @@ class DatetimeBlock(DatetimeLikeBlockMixin, Block):
     _can_hold_na = True
 
     def __init__(self, values, placement, fastpath=False, **kwargs):
-        if values.dtype != _NS_DTYPE:
+        if values.dtype != _NS_DTYPE and values.dtype.base != _NS_DTYPE:
+            # not datetime64 or datetime64tz
             values = conversion.ensure_datetime64ns(values)
 
         super(DatetimeBlock, self).__init__(values, fastpath=True,

pandas/tests/internals/test_external_block.py

@@ -5,12 +5,12 @@
 
 import pandas as pd
 from pandas.core.internals import (
-    Block, BlockManager, SingleBlockManager, NonConsolidatableMixIn)
+    BlockManager, SingleBlockManager, ExtensionBlock)
 
 import pytest
 
 
-class CustomBlock(NonConsolidatableMixIn, Block):
+class CustomBlock(ExtensionBlock):
 
     _holder = np.ndarray