REF: move common arithmetic ops from Integer/FloatingArray to NumericArray (#38290)

Author: jorisvandenbossche

pandas/core/arrays/floating.py

@@ -13,9 +13,7 @@
 from pandas.core.dtypes.common import (
     is_bool_dtype,
     is_datetime64_dtype,
-    is_float,
     is_float_dtype,
-    is_integer,
     is_integer_dtype,
     is_list_like,
     is_object_dtype,
@@ -28,7 +26,8 @@
 from pandas.core.ops import invalid_comparison
 from pandas.core.tools.numeric import to_numeric
 
-from .masked import BaseMaskedArray, BaseMaskedDtype
+from .masked import BaseMaskedDtype
+from .numeric import NumericArray
 
 if TYPE_CHECKING:
     import pyarrow
@@ -199,7 +198,7 @@ def coerce_to_array(
     return values, mask
 
 
-class FloatingArray(BaseMaskedArray):
+class FloatingArray(NumericArray):
     """
     Array of floating (optional missing) values.
 
@@ -478,71 +477,6 @@ def _maybe_mask_result(self, result, mask, other, op_name: str):
 
         return type(self)(result, mask, copy=False)
 
-    def _arith_method(self, other, op):
-        from pandas.arrays import IntegerArray
-
-        omask = None
-
-        if getattr(other, "ndim", 0) > 1:
-            raise NotImplementedError("can only perform ops with 1-d structures")
-
-        if isinstance(other, (IntegerArray, FloatingArray)):
-            other, omask = other._data, other._mask
-
-        elif is_list_like(other):
-            other = np.asarray(other)
-            if other.ndim > 1:
-                raise NotImplementedError("can only perform ops with 1-d structures")
-            if len(self) != len(other):
-                raise ValueError("Lengths must match")
-            if not (is_float_dtype(other) or is_integer_dtype(other)):
-                raise TypeError("can only perform ops with numeric values")
-
-        else:
-            if not (is_float(other) or is_integer(other) or other is libmissing.NA):
-                raise TypeError("can only perform ops with numeric values")
-
-        if omask is None:
-            mask = self._mask.copy()
-            if other is libmissing.NA:
-                mask |= True
-        else:
-            mask = self._mask | omask
-
-        if op.__name__ == "pow":
-            # 1 ** x is 1.
-            mask = np.where((self._data == 1) & ~self._mask, False, mask)
-            # x ** 0 is 1.
-            if omask is not None:
-                mask = np.where((other == 0) & ~omask, False, mask)
-            elif other is not libmissing.NA:
-                mask = np.where(other == 0, False, mask)
-
-        elif op.__name__ == "rpow":
-            # 1 ** x is 1.
-            if omask is not None:
-                mask = np.where((other == 1) & ~omask, False, mask)
-            elif other is not libmissing.NA:
-                mask = np.where(other == 1, False, mask)
-            # x ** 0 is 1.
-            mask = np.where((self._data == 0) & ~self._mask, False, mask)
-
-        if other is libmissing.NA:
-            result = np.ones_like(self._data)
-        else:
-            with np.errstate(all="ignore"):
-                result = op(self._data, other)
-
-        # divmod returns a tuple
-        if op.__name__ == "divmod":
-            div, mod = result
-            return (
-                self._maybe_mask_result(div, mask, other, "floordiv"),
-                self._maybe_mask_result(mod, mask, other, "mod"),
-            )
-
-        return self._maybe_mask_result(result, mask, other, op.__name__)
-
 
 _dtype_docstring = """
 An ExtensionDtype for {dtype} data.

pandas/core/arrays/integer.py

@@ -1,11 +1,10 @@
-from datetime import timedelta
 import numbers
 from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Type, Union
 import warnings
 
 import numpy as np
 
-from pandas._libs import Timedelta, iNaT, lib, missing as libmissing
+from pandas._libs import iNaT, lib, missing as libmissing
 from pandas._typing import ArrayLike, DtypeObj
 from pandas.compat.numpy import function as nv
 from pandas.util._decorators import cache_readonly
@@ -16,7 +15,6 @@
     is_datetime64_dtype,
     is_float,
     is_float_dtype,
-    is_integer,
     is_integer_dtype,
     is_list_like,
     is_object_dtype,
@@ -29,6 +27,7 @@
 from pandas.core.tools.numeric import to_numeric
 
 from .masked import BaseMaskedArray, BaseMaskedDtype
+from .numeric import NumericArray
 
 if TYPE_CHECKING:
     import pyarrow
@@ -263,7 +262,7 @@ def coerce_to_array(
     return values, mask
 
 
-class IntegerArray(BaseMaskedArray):
+class IntegerArray(NumericArray):
     """
     Array of integer (optional missing) values.
 
@@ -494,7 +493,7 @@ def _values_for_argsort(self) -> np.ndarray:
         return data
 
     def _cmp_method(self, other, op):
-        from pandas.core.arrays import BaseMaskedArray, BooleanArray
+        from pandas.core.arrays import BooleanArray
 
         mask = None
 
@@ -538,75 +537,6 @@ def _cmp_method(self, other, op):
 
         return BooleanArray(result, mask)
 
-    def _arith_method(self, other, op):
-        from pandas.core.arrays import FloatingArray
-
-        op_name = op.__name__
-        omask = None
-
-        if getattr(other, "ndim", 0) > 1:
-            raise NotImplementedError("can only perform ops with 1-d structures")
-
-        if isinstance(other, (IntegerArray, FloatingArray)):
-            other, omask = other._data, other._mask
-
-        elif is_list_like(other):
-            other = np.asarray(other)
-            if other.ndim > 1:
-                raise NotImplementedError("can only perform ops with 1-d structures")
-            if len(self) != len(other):
-                raise ValueError("Lengths must match")
-            if not (is_float_dtype(other) or is_integer_dtype(other)):
-                raise TypeError("can only perform ops with numeric values")
-
-        elif isinstance(other, (timedelta, np.timedelta64)):
-            other = Timedelta(other)
-
-        else:
-            if not (is_float(other) or is_integer(other) or other is libmissing.NA):
-                raise TypeError("can only perform ops with numeric values")
-
-        if omask is None:
-            mask = self._mask.copy()
-            if other is libmissing.NA:
-                mask |= True
-        else:
-            mask = self._mask | omask
-
-        if op_name == "pow":
-            # 1 ** x is 1.
-            mask = np.where((self._data == 1) & ~self._mask, False, mask)
-            # x ** 0 is 1.
-            if omask is not None:
-                mask = np.where((other == 0) & ~omask, False, mask)
-            elif other is not libmissing.NA:
-                mask = np.where(other == 0, False, mask)
-
-        elif op_name == "rpow":
-            # 1 ** x is 1.
-            if omask is not None:
-                mask = np.where((other == 1) & ~omask, False, mask)
-            elif other is not libmissing.NA:
-                mask = np.where(other == 1, False, mask)
-            # x ** 0 is 1.
-            mask = np.where((self._data == 0) & ~self._mask, False, mask)
-
-        if other is libmissing.NA:
-            result = np.ones_like(self._data)
-        else:
-            with np.errstate(all="ignore"):
-                result = op(self._data, other)
-
-        # divmod returns a tuple
-        if op_name == "divmod":
-            div, mod = result
-            return (
-                self._maybe_mask_result(div, mask, other, "floordiv"),
-                self._maybe_mask_result(mod, mask, other, "mod"),
-            )
-
-        return self._maybe_mask_result(result, mask, other, op_name)
-
     def sum(self, *, skipna=True, min_count=0, **kwargs):
         nv.validate_sum((), kwargs)
         return super()._reduce("sum", skipna=skipna, min_count=min_count)

pandas/core/arrays/numeric.py

@@ -0,0 +1,92 @@
+import datetime
+
+import numpy as np
+
+from pandas._libs import Timedelta, missing as libmissing
+from pandas.errors import AbstractMethodError
+
+from pandas.core.dtypes.common import (
+    is_float,
+    is_float_dtype,
+    is_integer,
+    is_integer_dtype,
+    is_list_like,
+)
+
+from .masked import BaseMaskedArray
+
+
+class NumericArray(BaseMaskedArray):
+    """
+    Base class for IntegerArray and FloatingArray.
+    """
+
+    def _maybe_mask_result(self, result, mask, other, op_name: str):
+        raise AbstractMethodError(self)
+
+    def _arith_method(self, other, op):
+        op_name = op.__name__
+        omask = None
+
+        if getattr(other, "ndim", 0) > 1:
+            raise NotImplementedError("can only perform ops with 1-d structures")
+
+        if isinstance(other, NumericArray):
+            other, omask = other._data, other._mask
+
+        elif is_list_like(other):
+            other = np.asarray(other)
+            if other.ndim > 1:
+                raise NotImplementedError("can only perform ops with 1-d structures")
+            if len(self) != len(other):
+                raise ValueError("Lengths must match")
+            if not (is_float_dtype(other) or is_integer_dtype(other)):
+                raise TypeError("can only perform ops with numeric values")
+
+        elif isinstance(other, (datetime.timedelta, np.timedelta64)):
+            other = Timedelta(other)
+
+        else:
+            if not (is_float(other) or is_integer(other) or other is libmissing.NA):
+                raise TypeError("can only perform ops with numeric values")
+
+        if omask is None:
+            mask = self._mask.copy()
+            if other is libmissing.NA:
+                mask |= True
+        else:
+            mask = self._mask | omask
+
+        if op_name == "pow":
+            # 1 ** x is 1.
+            mask = np.where((self._data == 1) & ~self._mask, False, mask)
+            # x ** 0 is 1.
+            if omask is not None:
+                mask = np.where((other == 0) & ~omask, False, mask)
+            elif other is not libmissing.NA:
+                mask = np.where(other == 0, False, mask)
+
+        elif op_name == "rpow":
+            # 1 ** x is 1.
+            if omask is not None:
+                mask = np.where((other == 1) & ~omask, False, mask)
+            elif other is not libmissing.NA:
+                mask = np.where(other == 1, False, mask)
+            # x ** 0 is 1.
+            mask = np.where((self._data == 0) & ~self._mask, False, mask)
+
+        if other is libmissing.NA:
+            result = np.ones_like(self._data)
+        else:
+            with np.errstate(all="ignore"):
+                result = op(self._data, other)
+
+        # divmod returns a tuple
+        if op_name == "divmod":
+            div, mod = result
+            return (
+                self._maybe_mask_result(div, mask, other, "floordiv"),
+                self._maybe_mask_result(mod, mask, other, "mod"),
+            )
+
+        return self._maybe_mask_result(result, mask, other, op_name)