BUG: treat None as NA in DataFrame arithmetic operations, #992

Author: wesm

pandas/core/frame.py

@@ -168,35 +168,54 @@
 #----------------------------------------------------------------------
 # Factory helper methods
 
-def _arith_method(func, name, default_axis='columns'):
+def _arith_method(op, name, default_axis='columns'):
+    def na_op(x, y):
+        try:
+            result = op(x, y)
+        except TypeError:
+            xrav = x.ravel()
+            result = np.empty(x.size, dtype=x.dtype)
+            if isinstance(y, np.ndarray):
+                yrav = y.ravel()
+                mask = notnull(xrav) & notnull(yrav)
+                result[mask] = op(xrav[mask], yrav[mask])
+            else:
+                mask = notnull(xrav)
+                result[mask] = op(xrav[mask], y)
+
+            np.putmask(result, -mask, np.nan)
+            result = result.reshape(x.shape)
+
+        return result
+
     @Appender(_arith_doc % name)
     def f(self, other, axis=default_axis, level=None, fill_value=None):
         if isinstance(other, DataFrame):    # Another DataFrame
-            return self._combine_frame(other, func, fill_value, level)
+            return self._combine_frame(other, na_op, fill_value, level)
         elif isinstance(other, Series):
-            return self._combine_series(other, func, fill_value, axis, level)
+            return self._combine_series(other, na_op, fill_value, axis, level)
         elif isinstance(other, (list, tuple)):
             if axis is not None and self._get_axis_name(axis) == 'index':
                 casted = Series(other, index=self.index)
             else:
                 casted = Series(other, index=self.columns)
-            return self._combine_series(casted, func, fill_value, axis, level)
+            return self._combine_series(casted, na_op, fill_value, axis, level)
         elif isinstance(other, np.ndarray):
             if other.ndim == 1:
                 if axis is not None and self._get_axis_name(axis) == 'index':
                     casted = Series(other, index=self.index)
                 else:
                     casted = Series(other, index=self.columns)
-                return self._combine_series(casted, func, fill_value,
+                return self._combine_series(casted, na_op, fill_value,
                                             axis, level)
             elif other.ndim == 2:
                 casted = DataFrame(other, index=self.index,
                                    columns=self.columns)
-                return self._combine_frame(casted, func, fill_value, level)
+                return self._combine_frame(casted, na_op, fill_value, level)
             else:  # pragma: no cover
                 raise ValueError("Bad argument shape")
         else:
-            return self._combine_const(other, func)
+            return self._combine_const(other, na_op)
 
     f.__name__ = name
 

pandas/core/series.py

@@ -48,14 +48,14 @@ def na_op(x, y):
         try:
             result = op(x, y)
         except TypeError:
+            result = np.empty(len(x), dtype=x.dtype)
             if isinstance(y, np.ndarray):
                 mask = notnull(x) & notnull(y)
-                result = np.empty(len(x), dtype=x.dtype)
                 result[mask] = op(x[mask], y[mask])
             else:
                 mask = notnull(x)
-                result = np.empty(len(x), dtype=x.dtype)
                 result[mask] = op(x[mask], y)
+            np.putmask(result, -mask, np.nan)
 
         return result
 

pandas/tests/test_frame.py

@@ -2136,6 +2136,30 @@ def test_operators(self):
         expected = self.frame2 * 2
         assert_frame_equal(added, expected)
 
+    def test_operators_none_as_na(self):
+        df = DataFrame({"col1": [2,5.0,123,None],
+                        "col2": [1,2,3,4]})
+
+        ops = [operator.add, operator.sub, operator.mul, operator.div]
+
+        for op in ops:
+            filled = df.fillna(np.nan)
+            result = op(df, 3)
+            expected = op(filled, 3)
+            expected[com.isnull(expected)] = None
+            assert_frame_equal(result, expected)
+
+            result = op(df, df)
+            expected = op(filled, filled)
+            expected[com.isnull(expected)] = None
+            assert_frame_equal(result, expected)
+
+            result = op(df, df.fillna(7))
+            assert_frame_equal(result, expected)
+
+            result = op(df.fillna(7), df)
+            assert_frame_equal(result, expected)
+
     def test_logical_operators(self):
         import operator
 
@@ -2154,16 +2178,21 @@ def _check_unary_op(op):
             assert_frame_equal(result, expected)
 
         df1 = {'a': {'a': True, 'b': False, 'c': False, 'd': True, 'e': True},
-               'b': {'a': False, 'b': True, 'c': False, 'd': False, 'e': False},
-               'c': {'a': False, 'b': False, 'c': True, 'd': False, 'e': False},
+               'b': {'a': False, 'b': True, 'c': False,
+                     'd': False, 'e': False},
+               'c': {'a': False, 'b': False, 'c': True,
+                     'd': False, 'e': False},
                'd': {'a': True, 'b': False, 'c': False, 'd': True, 'e': True},
                'e': {'a': True, 'b': False, 'c': False, 'd': True, 'e': True}}
 
         df2 = {'a': {'a': True, 'b': False, 'c': True, 'd': False, 'e': False},
-               'b': {'a': False, 'b': True, 'c': False, 'd': False, 'e': False},
+               'b': {'a': False, 'b': True, 'c': False,
+                     'd': False, 'e': False},
                'c': {'a': True, 'b': False, 'c': True, 'd': False, 'e': False},
-               'd': {'a': False, 'b': False, 'c': False, 'd': True, 'e': False},
-               'e': {'a': False, 'b': False, 'c': False, 'd': False, 'e': True}}
+               'd': {'a': False, 'b': False, 'c': False,
+                     'd': True, 'e': False},
+               'e': {'a': False, 'b': False, 'c': False,
+                     'd': False, 'e': True}}
 
         df1 = DataFrame(df1)
         df2 = DataFrame(df2)