Implement "decode" parameter in pull()

Implement `decode (bool)` parameter in `pull()`. Decode the JSON data from the server into dicts. Only applies with `stream=True`.

Signed-off-by: Antonio <[email protected]>

podman/domain/images_manager.py

@@ -14,6 +14,7 @@
 from podman.api.http_utils import encode_auth_header
 from podman.domain.images import Image
 from podman.domain.images_build import BuildMixin
+from podman.domain.json_stream import json_stream
 from podman.domain.manager import Manager
 from podman.domain.registry_data import RegistryData
 from podman.errors import APIError, ImageNotFound, PodmanError
@@ -323,6 +324,8 @@ def pull(
             auth_config (Mapping[str, str]) – Override the credentials that are found in the
                 config for this request. auth_config should contain the username and password
                 keys to be valid.
+            decode (bool): Decode the JSON data from the server into dicts.
+                Only applies with ``stream=True``
             platform (str) – Platform in the format os[/arch[/variant]]
             progress_bar (bool) - Display a progress bar with the image pull progress (uses
                 the compat endpoint). Default: False
@@ -404,7 +407,7 @@ def pull(
             return None
 
         if stream:
-            return response.iter_lines()
+            return self._stream_helper(response, decode=kwargs.get("decode"))
 
         for item in response.iter_lines():
             obj = json.loads(item)
@@ -541,3 +544,24 @@ def scp(
         response = self.client.post(f"/images/scp/{source}", params=params)
         response.raise_for_status()
         return response.json()
+
+    def _stream_helper(self, response, decode=False):
+        """Generator for data coming from a chunked-encoded HTTP response."""
+
+        if response.raw._fp.chunked:
+            if decode:
+                yield from json_stream(self._stream_helper(response, False))
+            else:
+                reader = response.raw
+                while not reader.closed:
+                    # this read call will block until we get a chunk
+                    data = reader.read(1)
+                    if not data:
+                        break
+                    if reader._fp.chunk_left:
+                        data += reader.read(reader._fp.chunk_left)
+                    yield data
+        else:
+            # Response isn't chunked, meaning we probably
+            # encountered an error immediately
+            yield self._result(response, json=decode)

podman/domain/json_stream.py

@@ -0,0 +1,74 @@
+import json
+import json.decoder
+
+from podman.errors import StreamParseError
+
+json_decoder = json.JSONDecoder()
+
+
+def stream_as_text(stream):
+    """
+    Given a stream of bytes or text, if any of the items in the stream
+    are bytes convert them to text.
+    This function can be removed once we return text streams
+    instead of byte streams.
+    """
+    for data in stream:
+        if not isinstance(data, str):
+            data = data.decode('utf-8', 'replace')
+        yield data
+
+
+def json_splitter(buffer):
+    """Attempt to parse a json object from a buffer. If there is at least one
+    object, return it and the rest of the buffer, otherwise return None.
+    """
+    buffer = buffer.strip()
+    try:
+        obj, index = json_decoder.raw_decode(buffer)
+        rest = buffer[json.decoder.WHITESPACE.match(buffer, index).end() :]
+        return obj, rest
+    except ValueError:
+        return None
+
+
+def json_stream(stream):
+    """Given a stream of text, return a stream of json objects.
+    This handles streams which are inconsistently buffered (some entries may
+    be newline delimited, and others are not).
+    """
+    return split_buffer(stream, json_splitter, json_decoder.decode)
+
+
+def line_splitter(buffer, separator='\n'):
+    index = buffer.find(str(separator))
+    if index == -1:
+        return None
+    return buffer[: index + 1], buffer[index + 1 :]
+
+
+def split_buffer(stream, splitter=None, decoder=lambda a: a):
+    """Given a generator which yields strings and a splitter function,
+    joins all input, splits on the separator and yields each chunk.
+    Unlike string.split(), each chunk includes the trailing
+    separator, except for the last one if none was found on the end
+    of the input.
+    """
+    splitter = splitter or line_splitter
+    buffered = ''
+
+    for data in stream_as_text(stream):
+        buffered += data
+        while True:
+            buffer_split = splitter(buffered)
+            if buffer_split is None:
+                break
+
+            item, buffered = buffer_split
+            yield item
+
+    if buffered:
+        try:
+            yield decoder(buffered)
+        except Exception as e:
+            raise StreamParseError(e) from e

podman/errors/__init__.py

@@ -21,6 +21,7 @@
     'NotFound',
     'NotFoundError',
     'PodmanError',
+    'StreamParseError',
 ]
 
 try:
@@ -32,6 +33,7 @@
         InvalidArgument,
         NotFound,
         PodmanError,
+        StreamParseError,
     )
 except ImportError:
     pass

podman/errors/exceptions.py

@@ -142,3 +142,8 @@ def __init__(
 
 class InvalidArgument(PodmanError):
     """Parameter to method/function was not valid."""
+
+
+class StreamParseError(RuntimeError):
+    def __init__(self, reason):
+        self.msg = reason