Make serializer interface streaming

Author: ntninja

datastore/__init__.py

@@ -15,7 +15,8 @@
 	"BinaryNullDatastore", "BinaryDictDatastore",
 	"ObjectNullDatastore", "ObjectDictDatastore",
 	"Query", "Cursor",
-	"SerializerAdapter",
+	"SerializerAdapter", "SerializingError",
+	"SerializeError", "ParseError",
 
 	"abc", "typing", "util"
 )
@@ -38,6 +39,9 @@
 
 # import core.serialize
 from .core.serialize import SerializerAdapter
+from .core.serialize import SerializingError
+from .core.serialize import SerializeError
+from .core.serialize import ParseError
 
 
 ### Exposed submodules ###

datastore/core/serialize.py

@@ -14,7 +14,35 @@ class util:  # noqa
 T_co = typing.TypeVar("T_co", covariant=True)
 
 
-#FIXME: This stuff should support streaming data to the maximum extent possible
+class SerializingError(ValueError):
+	"""Base class for all encoding and decoding related errors."""
+
+	def __init__(self, message, encoder_name):
+		self.encoder_name = encoder_name
+
+		super().__init__(message)
+
+
+class SerializeError(SerializingError):
+	"""Raised when encoding a Python object into a byte string has failed
+	due to some problem with the input data."""
+
+	def __init__(self, encoder_name, original):
+		self.original = original
+
+		msg = "Object serialization error: {}".format(original)
+		super().__init__(msg, encoder_name)
+
+
+class ParseError(SerializingError):
+	"""Raised when decoding a byte string to a Python object has failed due to
+	some problem with the input data."""
+
+	def __init__(self, encoder_name, original):
+		self.original = original
+
+		msg = "Object parsing error: {}".format(original)
+		super().__init__(msg, encoder_name)
 
 
 class Serializer(typing.Generic[T_co], metaclass=abc.ABCMeta):
@@ -24,13 +52,15 @@ class Serializer(typing.Generic[T_co], metaclass=abc.ABCMeta):
 
 
 	@abc.abstractmethod
-	def loads(self, value: bytes) -> typing.List[T_co]:
+	def parse(self, value: util.stream.ReceiveStream) \
+	    -> util.stream.ArbitraryReceiveChannel[T_co]:
 		"""returns deserialized `value`."""
 		pass
 
 
 	@abc.abstractmethod
-	def dumps(self, value: typing.List[T_co]) -> bytes:
+	def serialize(self, value: util.stream.ReceiveChannel[T_co]) \
+	    -> util.stream.ArbitraryReceiveStream:
 		"""returns serialized `value`."""
 		pass
 
@@ -87,8 +117,8 @@ async def get(self, key: key_.Key) -> util.stream.ReceiveChannel[T_co]:
 		key
 			Key naming the object to retrieve
 		"""
-		value = await (await self.child_datastore.get(key)).collect()  #FIXME
-		return util.stream.receive_channel_from(self.serializer.loads(value))
+		value = await self.child_datastore.get(key)
+		return util.stream.receive_channel_from(self.serializer.parse(value))
 
 
 	async def _put(self, key: key_.Key, value: util.stream.ReceiveChannel[T_co]) -> None:
@@ -104,8 +134,8 @@ async def _put(self, key: key_.Key, value: util.stream.ReceiveChannel[T_co]) ->
 		value
 			The object to store.
 		"""
-		value_items = await value.collect()  #FIXME
-		value_bytes = self.serializer.dumps(value_items)
+		value_items = util.stream.receive_channel_from(value)
+		value_bytes = self.serializer.serialize(value_items)
 		await self.child_datastore.put(key, value_bytes)
 
 
@@ -130,7 +160,10 @@ async def query(self, query: query_.Query) -> query_.Cursor:
 		result = cursor._iterable
 		cursor._iterable = []
 		for field in result:
-			cursor._iterable.extend(self.serializer.loads(field))
+			field_stream = util.stream.receive_stream_from(field)
+			parse_items = util.stream.receive_channel_from(self.serializer.parse(field_stream))
+			async for item in parse_items:
+				cursor._iterable.append(item)
 
 		return cursor
 

datastore/core/util/stream.py

@@ -555,6 +555,15 @@ async def collect(self) -> bytes:
 					break
 				value += chunk
 		return bytes(value)
+	
+	def __aiter__(self):
+		return self
+
+	async def __anext__(self):
+		value = await self.receive_some()
+		if len(value) < 1:
+			raise StopAsyncIteration
+		return value
 
 
 class TeeingReceiveStream(ReceiveStream):

datastore/serializer/json.py

@@ -1,54 +1,192 @@
+import codecs
 import json
 import typing
 
+import datastore
 import datastore.abc
 
 __all__ = ("Serializer", "PrettySerializer")
 
 T_co = typing.TypeVar("T_co", covariant=True)
 
 
-#FIXME: This stuff should support streaming data to the maximum extent possible
+#XXX: Switch to [`ijson`](https://pypi.org/project/ijson/) once it supports a
+#     non-blocking parsing mode.
+#     Upstream issue: https://github.com/ICRAR/ijson/issues/22 (“working on this”)
 
 
 class Serializer(datastore.abc.Serializer[T_co], typing.Generic[T_co]):
-	"""json wrapper serializer that gives the most compact representation possible.
-	"""
-	
-	__slots__ = ("encoding", "indent", "separators")
+	"""JSON parser that handles concatenated JSON messages and encodes to the
+	   most compact representation possible by default."""
 
+	__slots__ = ("encoding", "indent", "separators",
+	             "_buffer", "_decoder1", "_decoder2", "_lasterror")
+
 	encoding:   str
 	indent:     typing.Optional[typing.Union[int, str]]
 	separators: typing.Tuple[str, str]
 
-	
+	_buffer:    typing.List[str]
+	_decoder1:  codecs.IncrementalDecoder
+	_decoder2:  json.JSONDecoder
+	_lasterror: typing.Optional[Exception]
+
+
 	def __init__(self, *, encoding: str = "utf-8",
 	             indent: typing.Optional[typing.Union[int, str]] = None,
 	             separators: typing.Tuple[str, str] = (',', ':')):
 		self.encoding   = encoding
 		self.indent     = indent
 		self.separators = separators
+		
+		self._buffer    = []
+		self._decoder1  = codecs.getincrementaldecoder(encoding)()
+		self._decoder2  = json.JSONDecoder()
+		self._lasterror = None
+
+	def parse_partial(self, data: bytes) -> typing.Iterator[T_co]:
+		"""Incrementally decodes JSON data sets into Python objects.
+
+		Raises
+		------
+		~datastore.DecodingError
+		"""
+		try:
+			# Python 3 requires all JSON data to be a text string
+			lines = self._decoder1.decode(data, False).split("\n")
+
+			# Add first input line to last buffer line, if applicable, to
+			# handle cases where the JSON string has been chopped in half
+			# at the network level due to streaming
+			if len(self._buffer) > 0 and self._buffer[-1] is not None:
+				self._buffer[-1] += lines[0]
+				self._buffer.extend(lines[1:])
+			else:
+				self._buffer.extend(lines)
+		except UnicodeDecodeError as error:
+			raise datastore.ParseError("json", error) from error
+
+		# Process data buffer
+		index = 0
+		try:
+			# Process each line as separate buffer
+			#PERF: This way the `.lstrip()` call becomes almost always a NOP
+			#      even if it does return a different string it will only
+			#      have to allocate a new buffer for the currently processed
+			#      line.
+			while index < len(self._buffer):
+				while self._buffer[index]:
+					# Make sure buffer does not start with whitespace
+					#PERF: `.lstrip()` does not reallocate if the string does
+					#      not actually start with whitespace.
+					self._buffer[index] = self._buffer[index].lstrip()
+
+					# Handle case where the remainder of the line contained
+					# only whitespace
+					if not self._buffer[index]:
+						self._buffer[index] = None
+						continue
+
+					# Try decoding the partial data buffer and return results
+					# from this
+					data = self._buffer[index]
+					for index2 in range(index, len(self._buffer)):
+						# If decoding doesn't succeed with the currently
+						# selected buffer (very unlikely with our current
+						# class of input data) then retry with appending
+						# any other pending pieces of input data
+						# This will happen with JSON data that contains
+						# arbitrary new-lines: "{1:\n2,\n3:4}"
+						if index2 > index:
+							data += "\n" + self._buffer[index2]
+
+						try:
+							print(repr(data))
+							(obj, offset) = self._decoder2.raw_decode(data)
+						except ValueError:
+							# Treat error as fatal if we have already added
+							# the final buffer to the input
+							if (index2 + 1) == len(self._buffer):
+								raise
+						else:
+							index = index2
+							break
+
+					# Decoding succeeded – yield result and shorten buffer
+					yield obj
+					
+					if (len(data) - offset) > 0:
+						buffer_offset = len(self._buffer[index]) - len(data) + offset
+						self._buffer[index] = self._buffer[index][buffer_offset:]
+					else:
+						self._buffer[index] = None
+				index += 1
+		except ValueError as error:
+			# It is unfortunately not possible to reliably detect whether
+			# parsing ended because of an error *within* the JSON string, or
+			# an unexpected *end* of the JSON string.
+			# We therefor have to assume that any error that occurs here
+			# *might* be related to the JSON parser hitting EOF and therefor
+			# have to postpone error reporting until `parse_finalize` is
+			# called.
+			self._lasterror = error
+		finally:
+			# Remove all processed buffers
+			del self._buffer[0:index]
+
+	def parse_finalize(self) -> typing.Iterator[T_co]:
+		"""Raises errors for incomplete buffered data that could not be parsed
+		because the end of the input data has been reached.
+
+		Raises
+		------
+		~ipfshttpclient.exceptions.DecodingError
+
+		Returns
+		-------
+			tuple : Always empty
+		"""
+		try:
+			try:
+				# Raise exception for remaining bytes in bytes decoder
+				self._decoder1.decode(b"", True)
+			except UnicodeDecodeError as error:
+				raise datastore.ParseError("json", error) from error
+
+			# Late raise errors that looked like they could have been fixed if
+			# the caller had provided more data
+			if self._buffer:
+				raise datastore.ParseError("json", self._lasterror) from self._lasterror
+		finally:
+			# Reset state
+			self._buffer    = []
+			self._lasterror = None
+			self._decoder1.reset()
+
+		return ()
 
-	
-	def loads(self, value: bytes) -> typing.List[T_co]:
-		"""returns json deserialized `value`."""
-		return [json.loads(value)]  #FIXME: Broken if more than one object
-	
-	
-	def dumps(self, value: typing.Iterable[T_co]) -> bytes:
-		"""returns json serialized `value` (pretty-printed)."""
-		result = bytearray()
-		for item in value:
-			# We force `ensure_ascii=False` here to most compact encoding of
-			# having non-ascii characters be encoded by the given encoding
-			# rather then using JSON string "\uXXXX" notation.
-			# We also force `sort_keys=True` to ensure the output is
-			# reproducible between interpreter runs.
-			result += json.dumps(item, ensure_ascii=False, sort_keys=True,
-			                     indent=self.indent, separators=self.separators
-			).encode(self.encoding)
-		return bytes(result)
+	async def parse(self, source: datastore.abc.ReceiveStream) -> typing.Iterator[T_co]:
+		async for chunk in source:
+			for obj in self.parse_partial(chunk):
+				yield obj
+		for obj in self.parse_finalize():
+			yield obj
 
+	async def serialize(self, value: datastore.abc.ReceiveChannel[T_co]) \
+	      -> typing.Iterator[bytes]:
+		"""returns json serialized `value` (pretty-printed)."""
+		try:
+			async for obj in value:
+				# We force `ensure_ascii=False` here to most compact encoding of
+				# having non-ascii characters be encoded by the given encoding
+				# rather then using JSON string "\uXXXX" notation.
+				# We also force `sort_keys=True` to ensure the output is
+				# reproducible between interpreter runs.
+				yield json.dumps(obj, ensure_ascii=False, sort_keys=True,
+				                 indent=self.indent, separators=self.separators
+				).encode(self.encoding)
+		except (UnicodeEncodeError, TypeError) as error:
+			raise datastore.SerializeError("json", error) from error
 
 
 class PrettySerializer(Serializer[T_co], typing.Generic[T_co]):