Merge pull request #282 from vincentbernat/fix/getitem

Implement Record.__getitem__()
netbox-community · Aug 18, 2020 · 83ee65c · 83ee65c
2 parents d7754e8 + eddf4d2
commit 83ee65c
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 2 deletions.
diff --git a/pynetbox/core/response.py b/pynetbox/core/response.py
@@ -207,8 +207,8 @@ def __iter__(self):
             else:
                 yield i, cur_attr
 
-    def __getitem__(self, item):
-        return item
+    def __getitem__(self, k):
+        return dict(self)[k]
 
     def __str__(self):
         return getattr(self, "name", None) or getattr(self, "label", None) or ""

diff --git a/tests/unit/test_response.py b/tests/unit/test_response.py
@@ -11,6 +11,36 @@
 
 
 class RecordTestCase(unittest.TestCase):
+    def test_attribute_access(self):
+        test_values = {
+            "id": 123,
+            "units": 12,
+            "nested_dict": {"id": 222, "name": "bar"},
+            "int_list": [123, 321, 231],
+        }
+        test_obj = Record(test_values, None, None)
+        self.assertEqual(test_obj.id, 123)
+        self.assertEqual(test_obj.units, 12)
+        self.assertEqual(test_obj.nested_dict.name, "bar")
+        self.assertEqual(test_obj.int_list[1], 321)
+        with self.assertRaises(AttributeError) as _:
+            test_obj.nothing
+
+    def test_dict_access(self):
+        test_values = {
+            "id": 123,
+            "units": 12,
+            "nested_dict": {"id": 222, "name": "bar"},
+            "int_list": [123, 321, 231],
+        }
+        test_obj = Record(test_values, None, None)
+        self.assertEqual(test_obj["id"], 123)
+        self.assertEqual(test_obj["units"], 12)
+        self.assertEqual(test_obj["nested_dict"]["name"], "bar")
+        self.assertEqual(test_obj["int_list"][1], 321)
+        with self.assertRaises(KeyError) as _:
+            test_obj["nothing"]
+
     def test_serialize_list_of_records(self):
         test_values = {
             "id": 123,