ENH: Support ExtensionArray operators via a mixin

doc/source/whatsnew/v0.24.0.txt

@@ -10,6 +10,27 @@ New features
 
 .. _whatsnew_0240.enhancements.other:
 
+``ExtensionArray`` operator support
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+An ``ExtensionArray`` subclass consisting of objects that have arithmetic and
+comparison operators defined on the underlying objects can easily support 
+those operators on the ``ExtensionArray``, and therefore the operators 
+on ``Series`` built on those ``ExtensionArray`` classes will work as expected.
+
+Two new mixin classes, :class:`ExtensionScalarArithmeticMixin` and
+:class:`ExtensionScalarComparisonMixin`, support this capability.  
+If developing an ``ExtensionArray`` subclass, for example ``MyExtensionArray``,
+simply include ``ExtensionScalarArithmeticMixin`` and/or 
+``ExtensionScalarComparisonMixin`` as parent classes of ``MyExtensionArray``
+as follows:
+
+.. code-block:: python
+
+    class MyExtensionArray(ExtensionArray, ExtensionScalarArithmeticMixin,
+                           ExtensionScalarComparisonMixin):
+
+
 Other Enhancements
 ^^^^^^^^^^^^^^^^^^
 - :func:`to_datetime` now supports the ``%Z`` and ``%z`` directive when passed into ``format`` (:issue:`13486`)

pandas/api/extensions/__init__.py

@@ -3,5 +3,8 @@
                                   register_index_accessor,
                                   register_series_accessor)
 from pandas.core.algorithms import take  # noqa
-from pandas.core.arrays.base import ExtensionArray  # noqa
+from pandas.core.arrays.base import (ExtensionArray,    # noqa
+                                     ExtensionScalarArithmeticMixin,
+                                     ExtensionScalarComparisonMixin,
+                                     ExtensionScalarOpsMixin)
 from pandas.core.dtypes.dtypes import ExtensionDtype  # noqa

pandas/conftest.py

@@ -83,7 +83,8 @@ def observed(request):
                              '__mul__', '__rmul__',
                              '__floordiv__', '__rfloordiv__',
                              '__truediv__', '__rtruediv__',
-                             '__pow__', '__rpow__']
+                             '__pow__', '__rpow__',
+                             '__mod__', '__rmod__']
 if not PY3:
     _all_arithmetic_operators.extend(['__div__', '__rdiv__'])
 
@@ -96,6 +97,15 @@ def all_arithmetic_operators(request):
     return request.param
 
 
+@pytest.fixture(params=['__eq__', '__ne__', '__le__',
+                        '__lt__', '__ge__', '__gt__'])
+def all_compare_operators(request):
+    """
+    Fixture for dunder names for common compare operations
+    """
+    return request.param
+
+
 @pytest.fixture(params=[None, 'gzip', 'bz2', 'zip',
                         pytest.param('xz', marks=td.skip_if_no_lzma)])
 def compression(request):

pandas/core/arrays/__init__.py

@@ -1,2 +1,5 @@
-from .base import ExtensionArray  # noqa
+from .base import (ExtensionArray,    # noqa
+                   ExtensionScalarArithmeticMixin,
+                   ExtensionScalarComparisonMixin,
+                   ExtensionScalarOpsMixin)
 from .categorical import Categorical  # noqa

pandas/core/arrays/base.py

@@ -7,8 +7,15 @@
 """
 import numpy as np
 
+import operator
+
 from pandas.errors import AbstractMethodError
 from pandas.compat.numpy import function as nv
+from pandas.compat import set_function_name, PY3
+from pandas.core.dtypes.common import (
+    is_extension_array_dtype,
+    is_list_like)
+from pandas.core import ops
 
 _not_implemented_message = "{} does not implement {}."
 
@@ -610,3 +617,124 @@ def _ndarray_values(self):
         used for interacting with our indexers.
         """
         return np.array(self)
+
+
+class ExtensionScalarOpsMixin(object):
+    """
+    A base class for the mixins for different operators.
+    Can also be used to define an individual method for a specific
+    operator using the class method create_method()
+    """
+
+    @classmethod
+    def _create_method(cls, op, coerce_to_dtype=True):
+        """
+        A class method that returns a method that will correspond to an
+        operator for an ExtensionArray subclass, by dispatching to the
+        relevant operator defined on the individual elements of the
+        ExtensionArray.
+
+        Parameters
+        ----------
+        op:               An operator that takes arguments op(a, b)
+        coerce_to_dtype:  boolean indicating whether to attempt to convert
+                          the result to the underlying ExtensionArray dtype
+                          (default True)
+
+        Returns
+        -------
+        A method that can be bound to a method of a class
+
+        Example
+        -------
+        Given an ExtensionArray subclass called MyClass, use
+
+        >>> __add__ = ExtensionScalarOpsMixin.create_method(operator.add)
+
+        in the class definition of MyClass to create the operator
+        for addition.
+
+        """
+
+        op_name = ops._get_op_name(op, False)
+
+        def _binop(self, other):
+            def convert_values(parm):
+                if isinstance(parm, ExtensionArray) or is_list_like(parm):
+                    ovalues = parm
+                elif is_extension_array_dtype(parm):
+                    ovalues = parm.values
+                else:  # Assume its an object
+                    ovalues = [parm] * len(self)
+                return ovalues
+            lvalues = convert_values(self)
+            rvalues = convert_values(other)
+
+            # If the operator is not defined for the underlying objects,
+            # a TypeError should be raised
+            res = [op(a, b) for (a, b) in zip(lvalues, rvalues)]
+
+            if coerce_to_dtype:
+                try:
+                    res = self._from_sequence(res)
+                except TypeError:
+                    pass
+
+            return res
+
+        name = '__{name}__'.format(name=op_name)
+        return set_function_name(_binop, name, cls)
+
+
+class ExtensionScalarArithmeticMixin(ExtensionScalarOpsMixin):
+    """A mixin for defining the arithmetic operations on an ExtensionArray
+    class, where it assumed that the underlying objects have the operators
+    already defined.
+
+    Usage
+    ------
+    If you have defined a subclass MyClass(ExtensionArray), then
+    use MyClass(ExtensionArray, ExtensionScalarArithmeticMixin) to
+    get the arithmetic operators
+    """
+
+    __add__ = ExtensionScalarOpsMixin._create_method(operator.add)
+    __radd__ = ExtensionScalarOpsMixin._create_method(ops.radd)
+    __sub__ = ExtensionScalarOpsMixin._create_method(operator.sub)
+    __rsub__ = ExtensionScalarOpsMixin._create_method(ops.rsub)
+    __mul__ = ExtensionScalarOpsMixin._create_method(operator.mul)
+    __rmul__ = ExtensionScalarOpsMixin._create_method(ops.rmul)
+    __pow__ = ExtensionScalarOpsMixin._create_method(operator.pow)
+    __rpow__ = ExtensionScalarOpsMixin._create_method(ops.rpow)
+    __mod__ = ExtensionScalarOpsMixin._create_method(operator.mod)
+    __rmod__ = ExtensionScalarOpsMixin._create_method(ops.rmod)
+    __floordiv__ = ExtensionScalarOpsMixin._create_method(operator.floordiv)
+    __rfloordiv__ = ExtensionScalarOpsMixin._create_method(ops.rfloordiv)
+    __truediv__ = ExtensionScalarOpsMixin._create_method(operator.truediv)
+    __rtruediv__ = ExtensionScalarOpsMixin._create_method(ops.rtruediv)
+    if not PY3:
+        __div__ = ExtensionScalarOpsMixin._create_method(operator.div)
+        __rdiv__ = ExtensionScalarOpsMixin._create_method(ops.rdiv)
+
+    __divmod__ = ExtensionScalarOpsMixin._create_method(divmod)
+    __rdivmod__ = ExtensionScalarOpsMixin._create_method(ops.rdivmod)
+
+
+class ExtensionScalarComparisonMixin(ExtensionScalarOpsMixin):
+    """A mixin for defining the comparison operations on an ExtensionArray
+    class, where it assumed that the underlying objects have the operators
+    already defined.
+
+    Usage
+    ------
+    If you have defined a subclass MyClass(ExtensionArray), then
+    use MyClass(ExtensionArray, ExtensionComparisonMixin) to
+    get the comparison operators
+    """
+
+    __eq__ = ExtensionScalarOpsMixin._create_method(operator.eq, False)
+    __ne__ = ExtensionScalarOpsMixin._create_method(operator.ne, False)
+    __lt__ = ExtensionScalarOpsMixin._create_method(operator.lt, False)
+    __gt__ = ExtensionScalarOpsMixin._create_method(operator.gt, False)
+    __le__ = ExtensionScalarOpsMixin._create_method(operator.le, False)
+    __ge__ = ExtensionScalarOpsMixin._create_method(operator.ge, False)

pandas/core/indexes/base.py

@@ -2972,16 +2972,20 @@ def get_value(self, series, key):
         # use this, e.g. DatetimeIndex
         s = getattr(series, '_values', None)
         if isinstance(s, (ExtensionArray, Index)) and is_scalar(key):
-            # GH 20825
+            # GH 20882, 21257
             # Unify Index and ExtensionArray treatment
             # First try to convert the key to a location
-            # If that fails, see if key is an integer, and
+            # If that fails, raise a KeyError if an integer
+            # index, otherwise, see if key is an integer, and
             # try that
             try:
                 iloc = self.get_loc(key)
                 return s[iloc]
             except KeyError:
-                if is_integer(key):
+                if (len(self) > 0 and
+                        self.inferred_type in ['integer', 'boolean']):
+                    raise
+                elif is_integer(key):
                     return s[key]
 
         s = com._values_from_object(series)

pandas/core/ops.py

@@ -30,6 +30,7 @@
     is_bool_dtype,
     is_list_like,
     is_scalar,
+    is_extension_array_dtype,
     _ensure_object)
 from pandas.core.dtypes.cast import (
     maybe_upcast_putmask, find_common_type,
@@ -990,6 +991,20 @@ def _construct_divmod_result(left, result, index, name, dtype):
     )
 
 
+def dispatch_to_extension_op(op, left, right):
+    """
+    Assume that left is a Series backed by an ExtensionArray,
+    apply the operator defined by op.
+    """
+
+    # This will raise TypeError if the op is not defined on the ExtensionArray
+    res_values = op(left.values, right)
+
+    res_name = get_op_result_name(left, right)
+    return left._constructor(res_values, index=left.index,
+                             name=res_name)
+
+
 def _arith_method_SERIES(cls, op, special):
     """
     Wrapper function for Series arithmetic operations, to avoid
@@ -1058,6 +1073,9 @@ def wrapper(left, right):
             raise TypeError("{typ} cannot perform the operation "
                             "{op}".format(typ=type(left).__name__, op=str_rep))
 
+        elif is_extension_array_dtype(left):
+            return dispatch_to_extension_op(op, left, right)
+
         lvalues = left.values
         rvalues = right
         if isinstance(rvalues, ABCSeries):
@@ -1208,6 +1226,9 @@ def wrapper(self, other, axis=None):
             return self._constructor(res_values, index=self.index,
                                      name=res_name)
 
+        elif is_extension_array_dtype(self):
+            return dispatch_to_extension_op(op, self, other)
+
         elif isinstance(other, ABCSeries):
             # By this point we have checked that self._indexed_same(other)
             res_values = na_op(self.values, other.values)

pandas/core/series.py

@@ -2203,23 +2203,22 @@ def _binop(self, other, func, level=None, fill_value=None):
             result.name = None
         return result
 
-    def combine(self, other, func, fill_value=np.nan):
+    def combine(self, other, func, fill_value=None):
         """
         Perform elementwise binary operation on two Series using given function
         with optional fill value when an index is missing from one Series or
         the other
-
         Parameters
         ----------
         other : Series or scalar value
         func : function
             Function that takes two scalars as inputs and return a scalar
         fill_value : scalar value
-
+            The default specifies to use the appropriate NaN value for
+            the underlying dtype of the Series
         Returns
         -------
         result : Series
-
         Examples
         --------
         >>> s1 = Series([1, 2])
@@ -2228,26 +2227,36 @@ def combine(self, other, func, fill_value=np.nan):
         0    0
         1    2
         dtype: int64
-
         See Also
         --------
         Series.combine_first : Combine Series values, choosing the calling
             Series's values first
         """
+        self_is_ext = is_extension_array_dtype(self.values)
+        if fill_value is None:
+            fill_value = na_value_for_dtype(self.dtype, False)
+
         if isinstance(other, Series):
             new_index = self.index.union(other.index)
             new_name = ops.get_op_result_name(self, other)
-            new_values = np.empty(len(new_index), dtype=self.dtype)
-            for i, idx in enumerate(new_index):
+            new_values = []
+            for idx in new_index:
                 lv = self.get(idx, fill_value)
                 rv = other.get(idx, fill_value)
                 with np.errstate(all='ignore'):
-                    new_values[i] = func(lv, rv)
+                    new_values.append(func(lv, rv))
         else:
             new_index = self.index
             with np.errstate(all='ignore'):
-                new_values = func(self._values, other)
+                new_values = [func(lv, other) for lv in self._values]
             new_name = self.name
+
+        if self_is_ext and not is_categorical_dtype(self.values):
+            try:
+                new_values = self._values._from_sequence(new_values)
+            except TypeError:
+                pass
+
         return self._constructor(new_values, index=new_index, name=new_name)
 
     def combine_first(self, other):

pandas/tests/extension/base/__init__.py

@@ -47,6 +47,7 @@ class TestMyDtype(BaseDtypeTests):
 from .groupby import BaseGroupbyTests  # noqa
 from .interface import BaseInterfaceTests  # noqa
 from .methods import BaseMethodsTests  # noqa
+from .ops import BaseOpsTests  # noqa
 from .missing import BaseMissingTests  # noqa
 from .reshaping import BaseReshapingTests  # noqa
 from .setitem import BaseSetitemTests  # noqa

pandas/tests/extension/base/getitem.py

@@ -130,7 +130,7 @@ def test_get(self, data):
         expected = s.iloc[[0, 1]]
         self.assert_series_equal(result, expected)
 
-        assert s.get(-1) == s.iloc[-1]
+        assert s.get(-1) is None
         assert s.get(s.index.max() + 1) is None
 
         s = pd.Series(data[:6], index=list('abcdef'))