Merge pull request #9 from Neradoc/iterator-on-objects

Iterator on JSON objects

Author: tannewt

adafruit_json_stream.py

@@ -154,7 +154,7 @@ def __init__(self, stream):
         self.finish_char = ""
 
     def finish(self):
-        """Consume all of the characters for this list from the stream."""
+        """Consume all of the characters for this container from the stream."""
         if not self.done:
             if self.active_child:
                 self.active_child.finish()
@@ -163,7 +163,8 @@ def finish(self):
         self.done = True
 
     def as_object(self):
-        """Consume all of the characters for this list from the stream and return as an object."""
+        """Consume all of the characters for this container from the stream
+        and return as an object."""
         if self.has_read:
             raise BufferError("Object has already been partly read.")
 
@@ -207,10 +208,17 @@ class TransientObject(Transient):
     def __init__(self, stream):
         super().__init__(stream)
         self.finish_char = "}"
-        self.active_child_key = None
+        self.active_key = None
+
+    def finish(self):
+        """Consume all of the characters for this container from the stream."""
+        if self.active_key and not self.active_child:
+            self.done = self.data.fast_forward(",")
+            self.active_key = None
+        super().finish()
 
     def __getitem__(self, key):
-        if self.active_child and self.active_child_key == key:
+        if self.active_child and self.active_key == key:
             return self.active_child
 
         self.has_read = True
@@ -219,12 +227,16 @@ def __getitem__(self, key):
             self.active_child.finish()
             self.done = self.data.fast_forward(",")
             self.active_child = None
-            self.active_child_key = None
+            self.active_key = None
         if self.done:
             raise KeyError(key)
 
         while not self.done:
-            current_key = self.data.next_value(":")
+            if self.active_key:
+                current_key = self.active_key
+                self.active_key = None
+            else:
+                current_key = self.data.next_value(":")
             if current_key is None:
                 self.done = True
                 break
@@ -234,11 +246,47 @@ def __getitem__(self, key):
                     self.done = True
                 if isinstance(next_value, Transient):
                     self.active_child = next_value
-                    self.active_child_key = key
+                    self.active_key = key
                 return next_value
             self.done = self.data.fast_forward(",")
         raise KeyError(key)
 
+    def __iter__(self):
+        return self
+
+    def _next_key(self):
+        """Return the next item's key, without consuming the value."""
+        if self.active_key:
+            if self.active_child:
+                self.active_child.finish()
+                self.active_child = None
+            self.done = self.data.fast_forward(",")
+            self.active_key = None
+        if self.done:
+            raise StopIteration()
+
+        self.has_read = True
+
+        current_key = self.data.next_value(":")
+        if current_key is None:
+            self.done = True
+            raise StopIteration()
+
+        self.active_key = current_key
+        return current_key
+
+    def __next__(self):
+        return self._next_key()
+
+    def items(self):
+        """Return iterator in the dictionary’s items ((key, value) pairs)."""
+        try:
+            while not self.done:
+                key = self._next_key()
+                yield (key, self[key])
+        except StopIteration:
+            return
+
 
 def load(data_iter):
     """Returns an object to represent the top level of the given JSON stream."""

examples/json_stream_local_file_advanced.py

@@ -0,0 +1,84 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 Scott Shawcroft for Adafruit Industries
+#
+# SPDX-License-Identifier: Unlicense
+
+import sys
+import time
+
+import adafruit_json_stream as json_stream
+
+# import json_stream
+
+
+class FakeResponse:
+    def __init__(self, file):
+        self.file = file
+
+    def iter_content(self, chunk_size):
+        while True:
+            yield self.file.read(chunk_size)
+
+
+f = open(sys.argv[1], "rb")  # pylint: disable=consider-using-with
+obj = json_stream.load(FakeResponse(f).iter_content(32))
+
+
+def find_keys(haystack, keys):
+    """If we don't know the order in which the keys are,
+    go through all of them and pick the ones we want"""
+    out = {}
+    # iterate on the items of an object
+    for key in haystack:
+        if key in keys:
+            # retrieve the value only if needed
+            value = haystack[key]
+            # if it's a sub object, get it all
+            if hasattr(value, "as_object"):
+                value = value.as_object()
+            out[key] = value
+    return out
+
+
+months = [
+    "January",
+    "February",
+    "March",
+    "April",
+    "May",
+    "June",
+    "July",
+    "August",
+    "September",
+    "October",
+    "November",
+    "December",
+]
+
+
+def time_to_date(stamp):
+    tt = time.localtime(stamp)
+    month = months[tt.tm_mon]
+    return f"{tt.tm_mday:2d}th of {month}"
+
+
+def ftoc(temp):
+    return (temp - 32) * 5 / 9
+
+
+currently = obj["currently"]
+print("Currently:")
+print(" ", time_to_date(currently["time"]))
+print(" ", currently["icon"])
+
+# iterate on the content of a list
+for i, day in enumerate(obj["daily"]["data"]):
+    day_items = find_keys(day, ("time", "summary", "temperatureHigh"))
+    date = time_to_date(day_items["time"])
+    print(
+        f'On {date}: {day_items["summary"]},',
+        f'Max: {int(day_items["temperatureHigh"])}F',
+        f'({int(ftoc(day_items["temperatureHigh"]))}C)',
+    )
+
+    if i > 4:
+        break

tests/test_json_stream.py

@@ -685,3 +685,78 @@ def test_as_object_grabbing_multiple_subscriptable_levels_again_after_passed_rai
     assert next(dict_1["sub_list"]) == "a"
     with pytest.raises(KeyError, match="sub_dict"):
         dict_1["sub_dict"]["sub_dict_name"]
+
+
+def test_iterating_keys(dict_with_keys):
+    """Iterate through keys of a simple object."""
+
+    bytes_io_chunk = BytesChunkIO(dict_with_keys.encode())
+    stream = adafruit_json_stream.load(bytes_io_chunk)
+    output = list(stream)
+    assert output == ["field_1", "field_2", "field_3"]
+
+
+def test_iterating_keys_get(dict_with_keys):
+    """Iterate through keys of a simple object and get values."""
+
+    the_dict = json.loads(dict_with_keys)
+
+    bytes_io_chunk = BytesChunkIO(dict_with_keys.encode())
+    stream = adafruit_json_stream.load(bytes_io_chunk)
+    for key in stream:
+        value = stream[key]
+        assert value == the_dict[key]
+
+
+def test_iterating_items(dict_with_keys):
+    """Iterate through items of a simple object."""
+
+    bytes_io_chunk = BytesChunkIO(dict_with_keys.encode())
+    stream = adafruit_json_stream.load(bytes_io_chunk)
+    output = list(stream.items())
+    assert output == [("field_1", 1), ("field_2", 2), ("field_3", 3)]
+
+
+def test_iterating_keys_after_get(dict_with_keys):
+    """Iterate through keys of a simple object after an item has already been read."""
+
+    bytes_io_chunk = BytesChunkIO(dict_with_keys.encode())
+    stream = adafruit_json_stream.load(bytes_io_chunk)
+    assert stream["field_1"] == 1
+    output = list(stream)
+    assert output == ["field_2", "field_3"]
+
+
+def test_iterating_items_after_get(dict_with_keys):
+    """Iterate through items of a simple object after an item has already been read."""
+
+    bytes_io_chunk = BytesChunkIO(dict_with_keys.encode())
+    stream = adafruit_json_stream.load(bytes_io_chunk)
+    assert stream["field_1"] == 1
+    output = list(stream.items())
+    assert output == [("field_2", 2), ("field_3", 3)]
+
+
+def test_iterating_complex_dict(complex_dict):
+    """Mix iterating over items of objects in objects in arrays."""
+
+    names = ["one", "two", "three", "four"]
+    sub_values = [None, "two point one", "three point one", None]
+
+    stream = adafruit_json_stream.load(BytesChunkIO(complex_dict.encode()))
+
+    thing_num = 0
+    for (index, item) in enumerate(stream.items()):
+        key, a_list = item
+        assert key == f"list_{index+1}"
+        for thing in a_list:
+            assert thing["dict_name"] == names[thing_num]
+            for sub_key in thing["sub_dict"]:
+                # break after getting a key with or without the value
+                # (testing finish() called from the parent list)
+                if sub_key == "sub_dict_name":
+                    if thing_num in {1, 2}:
+                        value = thing["sub_dict"][sub_key]
+                        assert value == sub_values[thing_num]
+                    break
+            thing_num += 1