From ee07366a1ad1e70168a27776a6312aea9045c62d Mon Sep 17 00:00:00 2001
From: Jean-Baptiste Bayle <j2b.bayle@gmail.com>
Date: Mon, 8 Mar 2021 20:48:38 +0100
Subject: [PATCH] Use implicit broadcasting to ForEachMOSA
---
lisainstrument/containers.py | 50 ++++++++++++++++++++++++++++++++++++
tests/test_containers.py | 22 ++++++++++++++++
2 files changed, 72 insertions(+)
diff --git a/lisainstrument/containers.py b/lisainstrument/containers.py
index 7d1e29d..f8fb922 100644
--- a/lisainstrument/containers.py
+++ b/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)
diff --git a/tests/test_containers.py b/tests/test_containers.py
index 42ea21b..32abe71 100644
--- a/tests/test_containers.py
+++ b/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."""
--
GitLab