Modified named_arrays.AbstractFunctionArray.__getitem__() to suppor…

…t any instance of `named_arrays.AbstractArray`. (#105)
sun-data · Jan 15, 2025 · 6f0f10e · 6f0f10e
1 parent fa8a0f8
commit 6f0f10e
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 13 deletions.
diff --git a/named_arrays/_functions/functions.py b/named_arrays/_functions/functions.py
@@ -174,12 +174,14 @@ def _getitem(
         shape_inputs = inputs.shape
         shape_outputs = outputs.shape
 
-        if isinstance(item, na.AbstractFunctionArray):
-            if np.any(item.inputs != self.inputs):
-                raise ValueError("boolean advanced index does not have the same inputs as the array")
-
-            item_inputs = item.outputs
-            item_outputs = item.outputs
+        if isinstance(item, na.AbstractArray):
+            if isinstance(item, na.AbstractFunctionArray):
+                if np.any(item.inputs != self.inputs):
+                    raise ValueError("boolean advanced index does not have the same inputs as the array")
+                item_inputs = item.outputs
+                item_outputs = item.outputs
+            else:
+                item_inputs = item_outputs = item
 
             shape_item_inputs = item_inputs.shape
             shape_item_outputs = item_outputs.shape

diff --git a/named_arrays/_functions/tests/test_functions.py b/named_arrays/_functions/tests/test_functions.py
@@ -131,10 +131,15 @@ def test_length(self, array: na.AbstractFunctionArray):
                     outputs=na.ScalarArrayRange(0, 2, axis='y'),
                 )
             ),
+            na.ScalarLinearSpace(0, 1, axis='y', num=_num_y) > 0.5,
             na.FunctionArray(
                 inputs=na.ScalarLinearSpace(0, 1, axis='y', num=_num_y),
                 outputs=na.ScalarLinearSpace(0, 1, axis='y', num=_num_y) > 0.5,
             ),
+            na.FunctionArray(
+                inputs=na.ScalarLinearSpace(0, 2, axis='y', num=_num_y),
+                outputs=na.ScalarLinearSpace(0, 1, axis='y', num=_num_y) > 0.5,
+            ),
             na.FunctionArray(
                 inputs=na.ScalarLinearSpace(0, 1, axis='y', num=_num_y),
                 outputs=na.UniformUncertainScalarArray(
@@ -157,13 +162,15 @@ def test__getitem__(
                 array[item]
             return
 
-        if isinstance(item, na.AbstractFunctionArray):
-            if np.any(item.inputs != array.inputs):
-                with pytest.raises(ValueError):
-                    array[item]
-                return
-
-            item_outputs = item_inputs = item.outputs
+        if isinstance(item, na.AbstractArray):
+            if isinstance(item, na.AbstractFunctionArray):
+                if np.any(item.inputs != array.inputs):
+                    with pytest.raises(ValueError):
+                        array[item]
+                    return
+                item_outputs = item_inputs = item.outputs
+            else:
+                item_outputs = item_inputs = item
 
         elif isinstance(item, dict):
             item_inputs = {