8th day of python challenges 111-117

venv/lib/python3.6/site-packages/pandas/tests/extension/base/ops.py

@@ -0,0 +1,173 @@
+import operator
+
+import pytest
+
+import pandas as pd
+from pandas.core import ops
+
+from .base import BaseExtensionTests
+
+
+class BaseOpsUtil(BaseExtensionTests):
+    def get_op_from_name(self, op_name):
+        short_opname = op_name.strip("_")
+        try:
+            op = getattr(operator, short_opname)
+        except AttributeError:
+            # Assume it is the reverse operator
+            rop = getattr(operator, short_opname[1:])
+            op = lambda x, y: rop(y, x)
+
+        return op
+
+    def check_opname(self, s, op_name, other, exc=Exception):
+        op = self.get_op_from_name(op_name)
+
+        self._check_op(s, op, other, op_name, exc)
+
+    def _check_op(self, s, op, other, op_name, exc=NotImplementedError):
+        if exc is None:
+            result = op(s, other)
+            expected = s.combine(other, op)
+            self.assert_series_equal(result, expected)
+        else:
+            with pytest.raises(exc):
+                op(s, other)
+
+    def _check_divmod_op(self, s, op, other, exc=Exception):
+        # divmod has multiple return values, so check separately
+        if exc is None:
+            result_div, result_mod = op(s, other)
+            if op is divmod:
+                expected_div, expected_mod = s // other, s % other
+            else:
+                expected_div, expected_mod = other // s, other % s
+            self.assert_series_equal(result_div, expected_div)
+            self.assert_series_equal(result_mod, expected_mod)
+        else:
+            with pytest.raises(exc):
+                divmod(s, other)
+
+
+class BaseArithmeticOpsTests(BaseOpsUtil):
+    """Various Series and DataFrame arithmetic ops methods.
+
+    Subclasses supporting various ops should set the class variables
+    to indicate that they support ops of that kind
+
+    * series_scalar_exc = TypeError
+    * frame_scalar_exc = TypeError
+    * series_array_exc = TypeError
+    * divmod_exc = TypeError
+    """
+
+    series_scalar_exc = TypeError
+    frame_scalar_exc = TypeError
+    series_array_exc = TypeError
+    divmod_exc = TypeError
+
+    def test_arith_series_with_scalar(self, data, all_arithmetic_operators):
+        # series & scalar
+        op_name = all_arithmetic_operators
+        s = pd.Series(data)
+        self.check_opname(s, op_name, s.iloc[0], exc=self.series_scalar_exc)
+
+    @pytest.mark.xfail(run=False, reason="_reduce needs implementation")
+    def test_arith_frame_with_scalar(self, data, all_arithmetic_operators):
+        # frame & scalar
+        op_name = all_arithmetic_operators
+        df = pd.DataFrame({"A": data})
+        self.check_opname(df, op_name, data[0], exc=self.frame_scalar_exc)
+
+    def test_arith_series_with_array(self, data, all_arithmetic_operators):
+        # ndarray & other series
+        op_name = all_arithmetic_operators
+        s = pd.Series(data)
+        self.check_opname(
+            s, op_name, pd.Series([s.iloc[0]] * len(s)), exc=self.series_array_exc
+        )
+
+    def test_divmod(self, data):
+        s = pd.Series(data)
+        self._check_divmod_op(s, divmod, 1, exc=self.divmod_exc)
+        self._check_divmod_op(1, ops.rdivmod, s, exc=self.divmod_exc)
+
+    def test_divmod_series_array(self, data, data_for_twos):
+        s = pd.Series(data)
+        self._check_divmod_op(s, divmod, data)
+
+        other = data_for_twos
+        self._check_divmod_op(other, ops.rdivmod, s)
+
+        other = pd.Series(other)
+        self._check_divmod_op(other, ops.rdivmod, s)
+
+    def test_add_series_with_extension_array(self, data):
+        s = pd.Series(data)
+        result = s + data
+        expected = pd.Series(data + data)
+        self.assert_series_equal(result, expected)
+
+    def test_error(self, data, all_arithmetic_operators):
+        # invalid ops
+        op_name = all_arithmetic_operators
+        with pytest.raises(AttributeError):
+            getattr(data, op_name)
+
+    def test_direct_arith_with_series_returns_not_implemented(self, data):
+        # EAs should return NotImplemented for ops with Series.
+        # Pandas takes care of unboxing the series and calling the EA's op.
+        other = pd.Series(data)
+        if hasattr(data, "__add__"):
+            result = data.__add__(other)
+            assert result is NotImplemented
+        else:
+            raise pytest.skip(
+                "{} does not implement add".format(data.__class__.__name__)
+            )
+
+
+class BaseComparisonOpsTests(BaseOpsUtil):
+    """Various Series and DataFrame comparison ops methods."""
+
+    def _compare_other(self, s, data, op_name, other):
+        op = self.get_op_from_name(op_name)
+        if op_name == "__eq__":
+            assert getattr(data, op_name)(other) is NotImplemented
+            assert not op(s, other).all()
+        elif op_name == "__ne__":
+            assert getattr(data, op_name)(other) is NotImplemented
+            assert op(s, other).all()
+
+        else:
+
+            # array
+            assert getattr(data, op_name)(other) is NotImplemented
+
+            # series
+            s = pd.Series(data)
+            with pytest.raises(TypeError):
+                op(s, other)
+
+    def test_compare_scalar(self, data, all_compare_operators):
+        op_name = all_compare_operators
+        s = pd.Series(data)
+        self._compare_other(s, data, op_name, 0)
+
+    def test_compare_array(self, data, all_compare_operators):
+        op_name = all_compare_operators
+        s = pd.Series(data)
+        other = pd.Series([data[0]] * len(data))
+        self._compare_other(s, data, op_name, other)
+
+    def test_direct_arith_with_series_returns_not_implemented(self, data):
+        # EAs should return NotImplemented for ops with Series.
+        # Pandas takes care of unboxing the series and calling the EA's op.
+        other = pd.Series(data)
+        if hasattr(data, "__eq__"):
+            result = data.__eq__(other)
+            assert result is NotImplemented
+        else:
+            raise pytest.skip(
+                "{} does not implement __eq__".format(data.__class__.__name__)
+            )