speed optimizations in wxf serializer and zlib encoder

Author: crazyeng

wolframclient/serializers/wl.py

@@ -1,6 +1,6 @@
 from __future__ import absolute_import, print_function, unicode_literals
 
-from itertools import chain
+from itertools import chain, starmap
 
 from wolframclient.serializers.base import FormatSerializer
 from wolframclient.serializers.utils import py_encode_decimal, py_encode_text
@@ -9,16 +9,14 @@
 from wolframclient.utils.encoding import force_bytes, force_text
 
 
-def yield_with_separators(iterable, separator=b", ", first=None, last=None):
-    if first:
-        yield first
+def yield_with_separators(iterable, first, last, separator=b", "):
+    yield first
     for i, arg in enumerate(iterable):
         if i:
             yield separator
         for sub in arg:
             yield sub
-    if last:
-        yield last
+    yield last
 
 
 class WLSerializer(FormatSerializer):
@@ -60,16 +58,14 @@ def serialize_int(self, number):
         yield b"%i" % number
 
     def serialize_rule(self, lhs, rhs):
-        return yield_with_separators((lhs, rhs), separator=b" -> ")
+        return chain(lhs, (b" -> ",), rhs)
 
     def serialize_rule_delayed(self, lhs, rhs):
-        return yield_with_separators((lhs, rhs), separator=b" :> ")
+        return chain(lhs, (b" :> ",), rhs)
 
     def serialize_mapping(self, mapping, **opts):
         return yield_with_separators(
-            (self.serialize_rule(key, value) for key, value in mapping),
-            first=b"<|",
-            last=b"|>",
+            starmap(self.serialize_rule, mapping), first=b"<|", last=b"|>"
         )
 
     def serialize_association(self, mapping, **opts):

wolframclient/serializers/wxf.py

@@ -1,6 +1,6 @@
 from __future__ import absolute_import, print_function, unicode_literals
 
-from itertools import chain
+from itertools import chain, starmap
 
 from wolframclient.serializers.base import FormatSerializer
 from wolframclient.serializers.utils import py_encode_decimal, safe_len
@@ -20,7 +20,12 @@
 )
 from wolframclient.utils import six
 from wolframclient.utils.api import zlib
-from wolframclient.utils.encoding import force_bytes, force_text
+from wolframclient.utils.encoding import concatenate_bytes, force_bytes, force_text
+from wolframclient.utils.functional import partition
+
+
+def serialize_rule(key, value, sep=(WXF_CONSTANTS.Rule,)):
+    return chain(sep, key, value)
 
 
 def get_length(iterable, length=None):
@@ -37,6 +42,16 @@ def get_length(iterable, length=None):
     return iterable, len(iterable)
 
 
+def compress(data):
+
+    compressor = zlib.compressobj()
+
+    for token in map(compressor.compress, map(concatenate_bytes, partition(data, 100))):
+        yield token
+
+    yield compressor.flush()
+
+
 class WXFSerializer(FormatSerializer):
     """ Serialize python objects to WXF. """
 
@@ -46,25 +61,14 @@ def __init__(self, normalizer=None, compress=False, **opts):
 
     def generate_bytes(self, data):
 
-        yield WXF_VERSION
-
         if self.compress:
-            yield WXF_HEADER_COMPRESS
 
-        yield WXF_HEADER_SEPARATOR
+            return chain(
+                (WXF_VERSION, WXF_HEADER_COMPRESS, WXF_HEADER_SEPARATOR),
+                compress(self.encode(data)),
+            )
 
-        if self.compress:
-            compressor = zlib.compressobj()
-            if six.PY2:
-                for payload in self.encode(data):
-                    yield compressor.compress(six.binary_type(payload))
-            else:
-                for payload in self.encode(data):
-                    yield compressor.compress(payload)
-            yield compressor.flush()
-        else:
-            for payload in self.encode(data):
-                yield payload
+        return chain((WXF_VERSION, WXF_HEADER_SEPARATOR), self.encode(data))
 
     def serialize_symbol(self, name):
         yield WXF_CONSTANTS.Symbol
@@ -116,12 +120,9 @@ def serialize_string(self, string):
 
     def serialize_bytes(self, bytes, as_byte_array=not six.PY2):
         if as_byte_array:
-            yield WXF_CONSTANTS.BinaryString
-            yield varint_bytes(len(bytes))
-            yield bytes
+            return (WXF_CONSTANTS.BinaryString, varint_bytes(len(bytes)), bytes)
         else:
-            for token in self.serialize_string(force_text(bytes, encoding="iso8859-1")):
-                yield token
+            return self.serialize_string(force_text(bytes, encoding="iso8859-1"))
 
     def serialize_mapping(self, keyvalue, **opts):
         # the normalizer is always sending an generator key, value
@@ -130,9 +131,7 @@ def serialize_mapping(self, keyvalue, **opts):
 
         return chain(
             (WXF_CONSTANTS.Association, varint_bytes(length)),
-            chain.from_iterable(
-                chain((WXF_CONSTANTS.Rule,), key, value) for key, value in iterable
-            ),
+            chain.from_iterable(starmap(serialize_rule, iterable)),
         )
 
     def serialize_numeric_array(self, data, dimensions, wl_type):