Allow multiplication of ForEachObject with arrays

lisainstrument/containers.py

@@ -20,6 +20,10 @@ logger = logging.getLogger(__name__)
 class ForEachObject(abc.ABC):
     """Abstract class which represents a dictionary holding a value for each object."""
 
+    # Used to bypass Numpy's vectorization and use
+    # this class `__rmul__()` method for multiplication from both sides
+    __array_priority__ = 10000
+
     def __init__(self, values, concurrent=False):
         """Initialize an object with a value or a function of the mosa index.
 

tests/test_containers.py

@@ -195,6 +195,33 @@ def test_multiplication():
         result = object_2 * object_3
 
 
+def test_multiplication_array():
+    """Check that we can multiply a ForEachObject instance with an array."""
+
+    array = np.random.normal(size=10)
+    object_1 = ForEachAB({'A': 1, 'B': 2})
+    object_2 = ForEachAB({
+        'A': np.random.normal(size=10),
+        'B': np.random.normal(size=10),
+    })
+
+    result = object_1 * array
+    assert np.all(result['A'] == array)
+    assert np.all(result['B'] == 2 * array)
+
+    result = array * object_1
+    assert np.all(result['A'] == array)
+    assert np.all(result['B'] == 2 * array)
+
+    result = object_2 * array
+    assert np.all(result['A'] == object_2['A'] * array)
+    assert np.all(result['B'] == object_2['B'] * array)
+
+    result = array * object_2
+    assert np.all(result['A'] == object_2['A'] * array)
+    assert np.all(result['B'] == object_2['B'] * array)
+
+
 def test_floor_division():
     """Check that we can apply floor division on two ForEachObject subclasses, or a scalar."""