Fixes indexing issue raised by #130 and adds tests

aerosandbox/dynamics/point_mass/common_point_mass.py

@@ -185,7 +185,7 @@ def makeline(k, v):
             control_variables
         ])
 
-    def __getitem__(self, index: int):
+    def __getitem__(self, index: Union[int, slice]):
         """
         Indexes one item from each attribute of a Dynamics instance.
         Returns a new Dynamics instance of the same type.
@@ -198,18 +198,21 @@ def __getitem__(self, index: int):
 
         """
         l = len(self)
-        if index >= l or index < -l:
-            raise IndexError("Index is out of range!")
 
         def get_item_of_attribute(a):
-            try:
-                return a[index]
-            except TypeError as e:  # object is not subscriptable
+            if hasattr(a, "__len__") and hasattr(a, "__getitem__"):
+                if len(a) == 1:
+                    return a[0]
+                elif len(a) == l:
+                    return a[index]
+                else:
+                    try:
+                        return a[index]
+                    except IndexError as e:
+                        raise IndexError(f"A state variable could not be indexed; it has length {len(a)} while the"
+                                         f"parent has length {l}.")
+            else:
                 return a
-            except IndexError as e:  # index out of range
-                raise IndexError("A state variable could not be indexed, since the index is out of range!")
-            except NotImplementedError as e:
-                raise TypeError(f"Indices must be integers or slices, not {index.__class__.__name__}")
 
         new_instance = self.get_new_instance_with_state()
 

aerosandbox/dynamics/test_dynamics/test_indexing.py

@@ -0,0 +1,34 @@
+import aerosandbox as asb
+import aerosandbox.numpy as np
+import pytest
+
+def test_dyn_indexing():
+
+    # Test indexing of a simple Dynamics object
+    dyn = asb.DynamicsPointMass1DHorizontal(
+        mass_props=asb.MassProperties(
+            mass=1
+        ),
+        x_e=np.arange(10) ** 2,
+        u_e=2 * np.arange(10),
+    )
+
+    d0 = dyn[0]
+    d1 = dyn[1]
+    dn1 = dyn[-1]
+    dn2 = dyn[-2]
+    dslice = dyn[2:5]
+
+    assert d0.x_e == 0
+    assert d0.u_e == 0
+    assert d1.x_e == 1
+    assert d1.u_e == 2
+    assert dn1.x_e == 81
+    assert dn1.u_e == 18
+    assert dn2.x_e == 64
+    assert dn2.u_e == 16
+    assert all(dslice.x_e == [4, 9, 16])
+    assert all(dslice.u_e == [4, 6, 8])
+
+if __name__ == '__main__':
+    test_dyn_indexing()

aerosandbox/performance/operating_point.py

@@ -189,7 +189,7 @@ def makeline(k, v):
             state_variables,
         ])
 
-    def __getitem__(self, index: int) -> "OperatingPoint":
+    def __getitem__(self, index: Union[int, slice]):
         """
         Indexes one item from each attribute of an OperatingPoint instance.
         Returns a new OperatingPoint instance.
@@ -202,18 +202,21 @@ def __getitem__(self, index: int) -> "OperatingPoint":
 
         """
         l = len(self)
-        if index >= l or index < -l:
-            raise IndexError("Index is out of range!")
 
         def get_item_of_attribute(a):
-            try:
-                return a[index]
-            except TypeError as e:  # object is not subscriptable
+            if hasattr(a, "__len__") and hasattr(a, "__getitem__"):
+                if len(a) == 1:
+                    return a[0]
+                elif len(a) == l:
+                    return a[index]
+                else:
+                    try:
+                        return a[index]
+                    except IndexError as e:
+                        raise IndexError(f"A state variable could not be indexed; it has length {len(a)} while the"
+                                         f"parent has length {l}.")
+            else:
                 return a
-            except IndexError as e:  # index out of range
-                raise IndexError("A state variable could not be indexed, since the index is out of range!")
-            except NotImplementedError as e:
-                raise TypeError(f"Indices must be integers or slices, not {index.__class__.__name__}")
 
         new_instance = self.get_new_instance_with_state()
 

aerosandbox/performance/test_performance/test_indexing.py

@@ -0,0 +1,68 @@
+import aerosandbox as asb
+import aerosandbox.numpy as np
+import pytest
+
+
+def test_op_indexing():
+
+    op = asb.OperatingPoint(
+        atmosphere=asb.Atmosphere(altitude=np.linspace(0, 100)),
+        velocity=np.linspace(30, 70),
+        alpha=np.linspace(-5, 5),
+        beta=0,
+    )
+    o0 = op[0]
+    o1 = op[1]
+    on1 = op[-1]
+    on2 = op[-2]
+    oslice = op[2:5]
+
+    assert o0.atmosphere == op.atmosphere[0]
+    assert o0.alpha == op.alpha[0]
+    assert o0.velocity == op.velocity[0]
+
+    assert o1.atmosphere == op.atmosphere[1]
+    assert o1.alpha == op.alpha[1]
+    assert o1.velocity == op.velocity[1]
+
+    assert on1.atmosphere == op.atmosphere[-1]
+    assert on1.alpha == op.alpha[-1]
+    assert on1.velocity == op.velocity[-1]
+
+    assert on2.atmosphere == op.atmosphere[-2]
+    assert on2.alpha == op.alpha[-2]
+    assert on2.velocity == op.velocity[-2]
+
+    assert oslice.atmosphere == op.atmosphere[2:5]
+    assert np.all(oslice.alpha == op.alpha[2:5])
+    assert np.all(oslice.velocity == op.velocity[2:5])
+
+    # # Test indexing of a simple Dynamics object
+    # dyn = asb.DynamicsPointMass1DHorizontal(
+    #     mass_props=asb.MassProperties(
+    #         mass=1
+    #     ),
+    #     x_e=np.arange(10) ** 2,
+    #     u_e=2 * np.arange(10),
+    # )
+    #
+    # d0 = dyn[0]
+    # d1 = dyn[1]
+    # dn1 = dyn[-1]
+    # dn2 = dyn[-2]
+    # dslice = dyn[2:5]
+    #
+    # assert d0.x_e == 0
+    # assert d0.u_e == 0
+    # assert d1.x_e == 1
+    # assert d1.u_e == 2
+    # assert dn1.x_e == 81
+    # assert dn1.u_e == 18
+    # assert dn2.x_e == 64
+    # assert dn2.u_e == 16
+    # assert all(dslice.x_e == [4, 9, 16])
+    # assert all(dslice.u_e == [4, 6, 8])
+
+
+if __name__ == '__main__':
+    test_op_indexing()

aerosandbox/weights/mass_properties.py

@@ -131,18 +131,21 @@ def __getitem__(self, index) -> "MassProperties":
 
         """
         l = len(self)
-        if index >= l or index < -l:
-            raise IndexError("Index is out of range!")
 
         def get_item_of_attribute(a):
-            if np.isscalar(a):  # If NumPy says its a scalar, return it.
-                return a
-            try:
-                return a[index]
-            except TypeError:  # object is not subscriptable
+            if hasattr(a, "__len__") and hasattr(a, "__getitem__"):
+                if len(a) == 1:
+                    return a[0]
+                elif len(a) == l:
+                    return a[index]
+                else:
+                    try:
+                        return a[index]
+                    except IndexError as e:
+                        raise IndexError(f"A state variable could not be indexed; it has length {len(a)} while the"
+                                         f"parent has length {l}.")
+            else:
                 return a
-            except IndexError as e:  # index out of range
-                raise IndexError("A state variable could not be indexed, since the index is out of range!")
 
         inputs = {
             "mass": self.mass,