Add support for accessing nested items in BoxList using numpy-style t…

…uple indexing. (#266)
cdgriffith · Jun 12, 2024 · 51d8cfc · 51d8cfc
1 parent cc26a46
commit 51d8cfc
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/box/box_list.py b/box/box_list.py
@@ -65,6 +65,14 @@ def __getitem__(self, item):
             if len(list_pos.group()) == len(item):
                 return value
             return value.__getitem__(item[len(list_pos.group()) :].lstrip("."))
+        if isinstance(item, tuple):
+            result = self
+            for idx in item:
+                if isinstance(result, list):
+                    result = result[idx]
+                else:
+                    raise BoxTypeError(f"Cannot numpy-style indexing on {type(result).__name__}.")
+            return result
         return super().__getitem__(item)
 
     def __delitem__(self, key):

diff --git a/test/test_box_list.py b/test/test_box_list.py
@@ -35,6 +35,10 @@ def test_box_list(self):
         assert new_list[-1].item == 22
         new_list.append([{"bad_item": 33}])
         assert new_list[-1][0].bad_item == 33
+        new_list[-1].append([{"bad_item": 33}])
+        assert new_list[-1, -1, 0].bad_item == 33
+        bx = Box({0: {1: {2: {3: 3}}}, (0, 1, 2, 3): 4})
+        assert bx[0, 1, 2, 3] == 4
         assert repr(new_list).startswith("BoxList(")
         for x in new_list.to_list():
             assert not isinstance(x, (BoxList, Box))