using concatenate_bytes

Author: crazyeng

wolframclient/serializers/wxfencoder/streaming.py

@@ -1,8 +1,8 @@
 from __future__ import absolute_import, print_function, unicode_literals
 
-from wolframclient.utils import six
 from wolframclient.utils.api import zlib
-from wolframclient.utils.encoding import force_bytes
+from wolframclient.utils.decorators import decorate
+from wolframclient.utils.encoding import concatenate_bytes, force_bytes
 
 
 class ZipCompressedWriter(object):
@@ -32,6 +32,7 @@ class ExactSizeReader(object):
     def __init__(self, reader):
         self._reader = reader
 
+    @decorate(concatenate_bytes)
     def read(self, size=-1):
         """Read from an underlying readable object.
 
@@ -44,17 +45,16 @@ def read(self, size=-1):
         # Negative values read until EOF and 0 returns b''. Both remain unchanged.
         # Also a fast path when the requested amount of bytes is returned in one go.
         if size <= 0 or len(data) == size:
-            return data
-        # need an intermediary buffer
-        out_len = len(data)
-        data = six.BytesIO(data)
-        while out_len < size:
-            chunk = self._reader.read(size - out_len)
-            if chunk == b"":
-                raise EOFError("Not enough data to read.")
-            data.write(chunk)
-            out_len = out_len + len(chunk)
-        return data.getvalue()
+            yield data
+        else:
+            # need an intermediary buffer
+            out_len = len(data)
+            while out_len < size:
+                chunk = self._reader.read(size - out_len)
+                if chunk == b"":
+                    raise EOFError("Not enough data to read.")
+                yield chunk
+                out_len = out_len + len(chunk)
 
 
 class ZipCompressedReader(object):
@@ -70,6 +70,7 @@ def __init__(self, reader):
         self._compressor = zlib.decompressobj()
         self._reader = reader
 
+    @decorate(concatenate_bytes)
     def read(self, size=-1):
         """Read from a compressed stream of bytes and return the inflated byte sequence.
 
@@ -81,7 +82,7 @@ def read(self, size=-1):
             size = -1
         else:
             chunk_size = ZipCompressedReader.CHUNK_SIZE
-        out_data = six.BytesIO()
+
         out_len = 0
         while True:
             # first step find try to find some data to uncompress.
@@ -103,8 +104,7 @@ def read(self, size=-1):
             # increment output len.
             out_len = out_len + len(chunk)
             # write to buffer
-            out_data.write(chunk)
+            yield chunk
             # check requested size against output length.
             if size > 0 and out_len == size:
                 break
-        return out_data.getvalue()