fix #725 : binary operators +, -, *, / on sessions does not include non-LArray objects

Author: alixdamman

doc/source/changes/version_0_30.rst.inc

@@ -226,4 +226,6 @@ Fixes
   to an Excel Workbook (closes :issue:`713`).
 
 * fixed missing documentation of many functions in :ref:`Utility Functions <api-ufuncs>` section
-  of the API Reference (closes :issue:`698`).
+  of the API Reference (closes :issue:`698`).
+
+* fixed arithmetic operations between two sessions returning a nan value for each axis and group (closes :issue:`725`).

larray/core/session.py

@@ -820,7 +820,7 @@ def __len__(self):
         return len(self._objects)
 
     # binary operations are dispatched element-wise to all arrays (we consider Session as an array-like)
-    def _binop(opname):
+    def _binop(opname, arrays_only=True):
         opfullname = '__%s__' % opname
 
         def opmethod(self, other):
@@ -831,21 +831,24 @@ def opmethod(self, other):
             with np.errstate(call=_session_float_error_handler):
                 res = []
                 for name in all_keys:
-                    self_array = self.get(name, nan)
+                    self_item = self.get(name, nan)
                     other_operand = other.get(name, nan) if hasattr(other, 'get') else other
-                    try:
-                        res_array = getattr(self_array, opfullname)(other_operand)
-                    # TypeError for str arrays, ValueError for incompatible axes, ...
-                    except Exception:
-                        res_array = nan
-                    # this should only ever happen when self_array is a non Array (eg. nan)
-                    if res_array is NotImplemented:
+                    if arrays_only and not isinstance(self_item, LArray):
+                        res_item = self_item
+                    else:
                         try:
-                            res_array = getattr(other_operand, '__%s__' % inverseop(opname))(self_array)
+                            res_item = getattr(self_item, opfullname)(other_operand)
                         # TypeError for str arrays, ValueError for incompatible axes, ...
                         except Exception:
-                            res_array = nan
-                    res.append((name, res_array))
+                            res_item = nan
+                        # this should only ever happen when self_array is a non Array (eg. nan)
+                        if res_item is NotImplemented:
+                            try:
+                                res_item = getattr(other_operand, '__%s__' % inverseop(opname))(self_item)
+                            # TypeError for str arrays, ValueError for incompatible axes, ...
+                            except Exception:
+                                res_item = nan
+                    res.append((name, res_item))
             return Session(res)
         opmethod.__name__ = opfullname
         return opmethod
@@ -859,8 +862,8 @@ def opmethod(self, other):
     __truediv__ = _binop('truediv')
     __rtruediv__ = _binop('rtruediv')
 
-    __eq__ = _binop('eq')
-    __ne__ = _binop('ne')
+    __eq__ = _binop('eq', arrays_only=False)
+    __ne__ = _binop('ne', arrays_only=False)
 
     # element-wise method factory
     # unary operations are (also) dispatched element-wise to all arrays

larray/tests/test_session.py

@@ -9,7 +9,7 @@
 
 from larray.tests.common import assert_array_nan_equal, inputpath, tmp_path, meta
 from larray import (Session, Axis, LArray, Group, isnan, zeros_like, ndtest, ones_like,
-                    local_arrays, global_arrays, arrays)
+                    local_arrays, global_arrays, arrays, nan)
 from larray.util.misc import pickle
 
 try:
@@ -388,81 +388,105 @@ def test_element_equals(session):
 
 
 def test_eq(session):
-    sess = session.filter(kind=LArray)
-    expected = Session([('e', e), ('f', f), ('g', g)])
-    assert all([array.all() for array in (sess == expected).values()])
+    sess = session.filter(kind=(Axis, Group, LArray))
+    expected = Session([('b', b), ('b12', b12), ('a', a), ('a01', a01),
+                        ('e', e), ('g', g), ('f', f)])
+    assert all([item.all() if isinstance(item, LArray) else item
+                for item in (sess == expected).values()])
 
-    other = Session([('e', e), ('f', f)])
+    other = Session([('b', b), ('b12', b12), ('a', a), ('a01', a01), ('e', e), ('f', f)])
     res = sess == other
-    assert list(res.keys()) == ['e', 'g', 'f']
-    assert [arr.all() for arr in res.values()] == [True, False, True]
+    assert list(res.keys()) == ['b', 'b12', 'a', 'a01', 'e', 'g', 'f']
+    assert [item.all() if isinstance(item, LArray) else item
+            for item in res.values()] == [True, True, True, True, True, False, True]
 
     e2 = e.copy()
     e2.i[1, 1] = 42
-    other = Session([('e', e2), ('f', f)])
+    other = Session([('b', b), ('b12', b12), ('a', a), ('a01', a01), ('e', e2), ('f', f)])
     res = sess == other
-    assert [arr.all() for arr in res.values()] == [False, False, True]
+    assert [item.all() if isinstance(item, LArray) else item
+            for item in res.values()] == [True, True, True, True, False, False, True]
 
 
 def test_ne(session):
-    sess = session.filter(kind=LArray)
-    expected = Session([('e', e), ('f', f), ('g', g)])
-    assert ([(~array).all() for array in (sess != expected).values()])
+    sess = session.filter(kind=(Axis, Group, LArray))
+    expected = Session([('b', b), ('b12', b12), ('a', a), ('a01', a01),
+                        ('e', e), ('g', g), ('f', f)])
+    assert ([(~item).all() if isinstance(item, LArray) else not item
+             for item in (sess != expected).values()])
 
-    other = Session([('e', e), ('f', f)])
+    other = Session([('b', b), ('b12', b12), ('a', a), ('a01', a01), ('e', e), ('f', f)])
     res = sess != other
-    assert [(~arr).all() for arr in res.values()] == [True, False, True]
+    assert list(res.keys()) == ['b', 'b12', 'a', 'a01', 'e', 'g', 'f']
+    assert [(~item).all() if isinstance(item, LArray) else not item
+            for item in res.values()] == [True, True, True, True, True, False, True]
 
     e2 = e.copy()
     e2.i[1, 1] = 42
-    other = Session([('e', e2), ('f', f)])
+    other = Session([('b', b), ('b12', b12), ('a', a), ('a01', a01), ('e', e2), ('f', f)])
     res = sess != other
-    assert [(~arr).all() for arr in res.values()] == [False, False, True]
+    assert [(~item).all() if isinstance(item, LArray) else not item
+            for item in res.values()] == [True, True, True, True, False, False, True]
 
 
 def test_sub(session):
-    sess = session.filter(kind=LArray)
+    sess = session
 
     # session - session
     other = Session({'e': e - 1, 'f': ones_like(f)})
     diff = sess - other
     assert_array_nan_equal(diff['e'], np.full((2, 3), 1, dtype=np.int32))
     assert_array_nan_equal(diff['f'], f - ones_like(f))
     assert isnan(diff['g']).all()
+    assert diff.a is a
+    assert diff.a01 is a01
+    assert diff.c is c
 
     # session - scalar
     diff = sess - 2
     assert_array_nan_equal(diff['e'], e - 2)
     assert_array_nan_equal(diff['f'], f - 2)
     assert_array_nan_equal(diff['g'], g - 2)
+    assert diff.a is a
+    assert diff.a01 is a01
+    assert diff.c is c
 
     # session - dict(LArray and scalar)
     other = {'e': ones_like(e), 'f': 1}
     diff = sess - other
     assert_array_nan_equal(diff['e'], e - ones_like(e))
     assert_array_nan_equal(diff['f'], f - 1)
     assert isnan(diff['g']).all()
+    assert diff.a is a
+    assert diff.a01 is a01
+    assert diff.c is c
 
 
 def test_rsub(session):
-    sess = session.filter(kind=LArray)
+    sess = session
 
     # scalar - session
     diff = 2 - sess
     assert_array_nan_equal(diff['e'], 2 - e)
     assert_array_nan_equal(diff['f'], 2 - f)
     assert_array_nan_equal(diff['g'], 2 - g)
+    assert diff.a is a
+    assert diff.a01 is a01
+    assert diff.c is c
 
     # dict(LArray and scalar) - session
     other = {'e': ones_like(e), 'f': 1}
     diff = other - sess
     assert_array_nan_equal(diff['e'], ones_like(e) - e)
     assert_array_nan_equal(diff['f'], 1 - f)
     assert isnan(diff['g']).all()
+    assert diff.a is a
+    assert diff.a01 is a01
+    assert diff.c is c
 
 
 def test_div(session):
-    sess = session.filter(kind=LArray)
+    sess = session
     other = Session({'e': e - 1, 'f': f + 1})
 
     with pytest.warns(RuntimeWarning) as caught_warnings:
@@ -481,19 +505,25 @@ def test_div(session):
 
 
 def test_rdiv(session):
-    sess = session.filter(kind=LArray)
+    sess = session
 
     # scalar / session
     res = 2 / sess
     assert_array_nan_equal(res['e'], 2 / e)
     assert_array_nan_equal(res['f'], 2 / f)
     assert_array_nan_equal(res['g'], 2 / g)
+    assert res.a is a
+    assert res.a01 is a01
+    assert res.c is c
 
     # dict(LArray and scalar) - session
     other = {'e': e, 'f': f}
     res = other / sess
     assert_array_nan_equal(res['e'], e / e)
     assert_array_nan_equal(res['f'], f / f)
+    assert res.a is a
+    assert res.a01 is a01
+    assert res.c is c
 
 
 def test_pickle_roundtrip(session, meta):