calcex.Figure: enable use of sum()

ioex/calcex.py

@@ -74,7 +74,8 @@ class Figure(object):
         return not (self == other)
 
     def __add__(self, other):
-        assert isinstance(self, type(other))
+        if not isinstance(self, type(other)):
+            raise NotImplementedError('{!r} + {!r}'.format(self, other))
         assert not self.value is None
         assert not other.value is None
         if self.unit != other.unit:
@@ -82,8 +83,16 @@ class Figure(object):
         else:
             return type(self)(value=self.value + other.value, unit=self.unit)
 
+    def __radd__(self, other):
+        """ enables use of sum() """
+        if isinstance(other, int) and other == 0:
+            return copy.deepcopy(self)
+        else:
+            raise NotImplementedError('{!r} + {!r}'.format(other, self))
+
     def __sub__(self, other):
-        assert isinstance(self, type(other))
+        if not isinstance(self, type(other)):
+            raise NotImplementedError('{!r} - {!r}'.format(self, other))
         assert not self.value is None
         assert not other.value is None
         if self.unit != other.unit:

tests/calcex/test_figure.py

@@ -155,6 +155,21 @@ def test_add_persistent():
     assert Figure([1, 2], ['m']) == s
 
 
+def test_add_sum():
+    s = sum([Figure(1, 'm'), Figure(2, 'm'), Figure(3, 'm')])
+    assert Figure(6, 'm') == s
+
+
+@pytest.mark.parametrize(('a', 'b'), [
+    [Figure(1, 'm'), 0],
+    [Figure(1, 'm'), 'test'],
+    [1, Figure(1, 'm')],
+])
+def test_add_not_implemented(a, b):
+    with pytest.raises(NotImplementedError):
+        a + b
+
+
 @pytest.mark.parametrize(('a', 'b', 'expected_sum'), [
     [Figure(1, 'm'), Figure(2, 'm'), Figure(-1, 'm')],
     [Figure(-2, 'l'), Figure(-4, 'l'), Figure(2, 'l')],