Allow multiple subscriptable levels twice

Author: justmobilize

adafruit_json_stream.py

@@ -174,14 +174,19 @@ def __init__(self, stream):
         super().__init__(stream)
         self.finish_char = "}"
         self.start_char = "{"
+        self.active_child_key = None
 
     def __getitem__(self, key):
+        if self.active_child and self.active_child_key == key:
+            return self.active_child
+
         self.has_read = True
 
         if self.active_child:
             self.active_child.finish()
             self.done = self.data.fast_forward(",")
             self.active_child = None
+            self.active_child_key = None
         if self.done:
             raise KeyError(key)
 
@@ -196,6 +201,7 @@ def __getitem__(self, key):
                     self.done = True
                 if isinstance(next_value, Transient):
                     self.active_child = next_value
+                    self.active_child_key = key
                 return next_value
             self.done = self.data.fast_forward(",")
         raise KeyError(key)

tests/test_json_stream.py

@@ -481,8 +481,8 @@ def test_as_object_that_is_partially_read_raises(complex_dict):
         dict_1.as_object()
 
 
-def test_as_object_grabbing_multiple_subscriptable_levels_twice_raises(complex_dict):
-    """Test loading a complex dict and grabbing multiple subscriptable levels twice raises."""
+def test_as_object_grabbing_multiple_subscriptable_levels_twice(complex_dict):
+    """Test loading a complex dict and grabbing multiple subscriptable levels twice."""
 
     assert json.loads(complex_dict)
 
@@ -491,5 +491,23 @@ def test_as_object_grabbing_multiple_subscriptable_levels_twice_raises(complex_d
     list_1 = stream["list_1"]
     dict_1 = next(list_1)
     assert dict_1["sub_dict"]["sub_dict_id"] == 1.1
+    assert dict_1["sub_dict"]["sub_dict_name"] == "one point one"
+
+
+def test_as_object_grabbing_multiple_subscriptable_levels_again_after_passed_raises(
+    complex_dict,
+):
+    """
+    Test loading a complex dict and grabbing multiple subscriptable levels after passing it raises.
+    """
+
+    assert json.loads(complex_dict)
+
+    stream = adafruit_json_stream.load(BytesChunkIO(complex_dict.encode()))
+
+    list_1 = stream["list_1"]
+    dict_1 = next(list_1)
+    assert dict_1["sub_dict"]["sub_dict_id"] == 1.1
+    assert next(dict_1["sub_list"]) == "a"
     with pytest.raises(KeyError, match="sub_dict"):
         dict_1["sub_dict"]["sub_dict_name"]