Use implicit broadcasting to ForEachMOSA

lisainstrument/containers.py

@@ -240,6 +240,31 @@ class ForEachSC(ForEachObject):
         """Return a ForEachMOSA instance by sharing the spacecraft values on both MOSAs."""
         return ForEachMOSA(lambda mosa: self[ForEachMOSA.sc(mosa)])
 
+    def __add__(self, other):
+        if isinstance(other, ForEachMOSA):
+            return self.for_each_mosa() + other
+        return super().__add__(other)
+
+    def __sub__(self, other):
+        if isinstance(other, ForEachMOSA):
+            return self.for_each_mosa() - other
+        return super().__sub__(other)
+
+    def __mul__(self, other):
+        if isinstance(other, ForEachMOSA):
+            return self.for_each_mosa() * other
+        return super().__mul__(other)
+
+    def __floordiv__(self, other):
+        if isinstance(other, ForEachMOSA):
+            return self.for_each_mosa() // other
+        return super().__floordiv__(other)
+
+    def __truediv__(self, other):
+        if isinstance(other, ForEachMOSA):
+            return self.for_each_mosa() / other
+        return super().__truediv__(other)
+
 
 class ForEachMOSA(ForEachObject):
     """Represents a dictionary of values for each moveable optical subassembly (MOSA)."""
@@ -283,3 +308,28 @@ class ForEachMOSA(ForEachObject):
     def adjacent(self):
         """Return a ForEachMOSA instance for adjacent MOSAs."""
         return ForEachMOSA(lambda mosa: self[ForEachMOSA.adjacent_mosa(mosa)])
+
+    def __add__(self, other):
+        if isinstance(other, ForEachSC):
+            return self + other.for_each_mosa()
+        return super().__add__(other)
+
+    def __sub__(self, other):
+        if isinstance(other, ForEachSC):
+            return self - other.for_each_mosa()
+        return super().__sub__(other)
+
+    def __mul__(self, other):
+        if isinstance(other, ForEachSC):
+            return self * other.for_each_mosa()
+        return super().__mul__(other)
+
+    def __floordiv__(self, other):
+        if isinstance(other, ForEachSC):
+            return self // other.for_each_mosa()
+        return super().__floordiv__(other)
+
+    def __truediv__(self, other):
+        if isinstance(other, ForEachSC):
+            return self / other.for_each_mosa()
+        return super().__truediv__(other)

tests/test_containers.py

@@ -368,6 +368,28 @@ def test_foreachsc_to_foreachmosa():
     assert my_mosa['32'] == 30
 
 
+def test_auto_foreachsc_to_foreachmosa():
+    """Test that ForEachSC turn automatically into ForEachMOSA during operations."""
+
+    my_sc = ForEachSC(int)
+    my_mosa = ForEachMOSA(int)
+
+    my_add_1 = my_sc + my_mosa
+    my_add_2 = my_mosa + my_sc
+    assert my_add_1 == my_sc.for_each_mosa() + my_mosa
+    assert my_add_1 == my_add_2
+
+    my_sub_1 = my_sc - my_mosa
+    my_sub_2 = my_mosa - my_sc
+    assert my_sub_1 == my_sc.for_each_mosa() - my_mosa
+    assert my_sub_1 == -my_sub_2
+
+    my_mult_1 = my_sc * my_mosa
+    my_mult_2 = my_mosa * my_sc
+    assert my_mult_1 == my_sc.for_each_mosa() * my_mosa
+    assert my_mult_1 == my_mult_2
+
+
 def test_foreachmosa_distant():
     """Test that one can generate a ForEachMOSA instant for distant MOSAs."""