[azure-core/corehttp] Add JSONL and SSE streaming support

sdk/core/azure-core/CHANGELOG.md

@@ -1,5 +1,11 @@
 # Release History
 
+## 1.43.0b1 (Unreleased)
+
+### Features Added
+
+- Added streaming support via the `azure.core.streaming` module for JSONL and Server-Sent Events (SSE), including the `Stream`/`AsyncStream` iterators and the `JSONLDecoder`/`AsyncJSONLDecoder` and `SSEDecoder`/`AsyncSSEDecoder` decoders. #38806
+
 ## 1.42.0 (Unreleased)
 
 ### Features Added

sdk/core/azure-core/README.md

@@ -229,6 +229,27 @@ foo = Foo(
 )
 ```
 
+#### Streaming
+
+`azure.core` provides a stream-agnostic `Stream` iterator for consuming streaming responses. Currently, JSON Lines (JSONL) streaming is supported via the `JSONLDecoder`. Pass the streamed response together with a decoder and a `deserialization_callback` that receives the response and each decoded event:
+
+```python
+from azure.core.rest import HttpRequest, HttpResponse
+from azure.core.streaming import Stream, JSONLDecoder
+
+request = HttpRequest("GET", "https://example.com/stream")
+response = client.send_request(request, stream=True)
+
+def deserialize(response, event):
+    return event  # or deserialize into a model
+
+with Stream(response=response, decoder=JSONLDecoder(), deserialization_callback=deserialize) as stream:
+    for item in stream:
+        print(item)
+```
+
+An asynchronous equivalent is available using `AsyncStream` and `AsyncJSONLDecoder`.
+
 ## Logging
 
 Azure libraries follow the guidance of Python's standard [logging](https://docs.python.org/3/library/logging.html) module. By following the Python documentation on logging, you should be able to configure logging for Azure libraries effectively.

sdk/core/azure-core/azure/core/_version.py

@@ -9,4 +9,4 @@
 # regenerated.
 # --------------------------------------------------------------------------
 
-VERSION = "1.42.0"
+VERSION = "1.43.0b1"

sdk/core/azure-core/azure/core/streaming/__init__.py

@@ -0,0 +1,37 @@
+# --------------------------------------------------------------------------
+#
+# Copyright (c) Microsoft Corporation. All rights reserved.
+#
+# The MIT License (MIT)
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the ""Software""), to
+# deal in the Software without restriction, including without limitation the
+# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+# sell copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+# IN THE SOFTWARE.
+#
+# --------------------------------------------------------------------------
+
+from ._stream import Stream, AsyncStream
+from ._jsonl import JSONLEvent
+from ._sse import ServerSentEvent
+
+
+__all__ = [
+    "Stream",
+    "AsyncStream",
+    "JSONLEvent",
+    "ServerSentEvent",
+]

sdk/core/azure-core/azure/core/streaming/_decoders.py

@@ -0,0 +1,63 @@
+# --------------------------------------------------------------------------
+#
+# Copyright (c) Microsoft Corporation. All rights reserved.
+#
+# The MIT License (MIT)
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the ""Software""), to
+# deal in the Software without restriction, including without limitation the
+# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+# sell copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+# IN THE SOFTWARE.
+#
+# --------------------------------------------------------------------------
+
+from typing import Iterator, AsyncIterator, Protocol
+
+from typing_extensions import runtime_checkable, TypeVar
+
+
+T_co = TypeVar("T_co", covariant=True)
+
+
+@runtime_checkable
+class StreamDecoder(Protocol[T_co]):
+    """Protocol for stream decoders."""
+
+    def iter_events(self, iter_bytes: Iterator[bytes]) -> Iterator[T_co]:
+        """Iterate over events from a byte iterator.
+
+        :param iter_bytes: An iterator of byte chunks.
+        :type iter_bytes: Iterator[bytes]
+        :return: An iterator of decoded data.
+        :rtype: Iterator[DecodedType_co]
+        """
+        ...
+
+
+@runtime_checkable
+class AsyncStreamDecoder(Protocol[T_co]):
+    """Protocol for async stream decoders."""
+
+    # Why this isn't async def: https://mypy.readthedocs.io/en/stable/more_types.html#asynchronous-iterators
+    def aiter_events(self, iter_bytes: AsyncIterator[bytes]) -> AsyncIterator[T_co]:
+        """Asynchronously iterate over events from a byte iterator.
+
+        :param iter_bytes: An asynchronous iterator of byte chunks.
+        :type iter_bytes: AsyncIterator[bytes]
+        :return: An asynchronous iterator of decoded data.
+        :rtype: AsyncIterator[DecodedType_co]
+        """
+        ...

sdk/core/azure-core/azure/core/streaming/_jsonl.py

@@ -0,0 +1,144 @@
+# --------------------------------------------------------------------------
+#
+# Copyright (c) Microsoft Corporation. All rights reserved.
+#
+# The MIT License (MIT)
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the ""Software""), to
+# deal in the Software without restriction, including without limitation the
+# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+# sell copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+# IN THE SOFTWARE.
+#
+# --------------------------------------------------------------------------
+
+import codecs
+import json
+from contextlib import aclosing
+from typing import Iterator, AsyncIterator, Any, Optional, cast
+
+
+class JSONLEvent:
+    """A single JSON Lines (JSONL) event.
+
+    :ivar data: The raw JSONL record.
+    :vartype data: str or None
+    """
+
+    def __init__(
+        self,
+        *,
+        data: Optional[str] = None,
+    ) -> None:
+        self.data = data
+
+    def json(self) -> Any:
+        """Parse the event data as JSON.
+
+        :return: The parsed JSON value.
+        :rtype: Any
+        """
+        return json.loads(cast(str, self.data))
+
+
+def iter_lines(iter_bytes: Iterator[bytes]) -> Iterator[str]:
+    """Iterate over lines from a byte iterator.
+
+    :param iter_bytes: An iterator of byte chunks.
+    :type iter_bytes: Iterator[bytes]
+    :rtype: Iterator[str]
+    :return: An iterator of lines.
+    """
+    decoder = codecs.getincrementaldecoder("utf-8")()
+
+    # Split only on "\n" (tolerating "\r\n") rather than using str.splitlines(),
+    # which would also break on other Unicode boundaries (\v, \f, \x1c-\x1e, \x85,
+    # \u2028, \u2029) that are valid inside a JSONL record's string value.
+    decoded = ""
+    for chunk in iter_bytes:
+        decoded += decoder.decode(chunk)
+        if decoded:
+            decoded_lines = [line[:-1] if line.endswith("\r") else line for line in decoded.split("\n")]
+            yield from decoded_lines[:-1]
+            decoded = decoded_lines[-1]
+
+    decoded += decoder.decode(b"", final=True)
+    if decoded:
+        yield decoded[:-1] if decoded.endswith("\r") else decoded
+
+
+async def aiter_lines(iter_bytes: AsyncIterator[bytes]) -> AsyncIterator[str]:
+    """Iterate over lines from a byte iterator.
+
+    :param iter_bytes: An iterator of byte chunks.
+    :type iter_bytes: Iterator[bytes]
+    :rtype: Iterator[str]
+    :return: An iterator of lines.
+    """
+    decoder = codecs.getincrementaldecoder("utf-8")()
+
+    # Split only on "\n" (tolerating "\r\n") rather than using str.splitlines(),
+    # which would also break on other Unicode boundaries (\v, \f, \x1c-\x1e, \x85,
+    # \u2028, \u2029) that are valid inside a JSONL record's string value.
+    decoded = ""
+    try:
+        async for chunk in iter_bytes:
+            decoded += decoder.decode(chunk)
+            if decoded:
+                decoded_lines = [line[:-1] if line.endswith("\r") else line for line in decoded.split("\n")]
+                for line in decoded_lines[:-1]:
+                    yield line
+                decoded = decoded_lines[-1]
+    finally:
+        aclose = getattr(iter_bytes, "aclose", None)
+        if aclose is not None:
+            await aclose()
+
+    decoded += decoder.decode(b"", final=True)
+    if decoded:
+        yield decoded[:-1] if decoded.endswith("\r") else decoded
+
+
+class JSONLDecoder:
+    """Decoder for JSON Lines (JSONL) format. https://jsonlines.org/"""
+
+    def iter_events(self, iter_bytes: Iterator[bytes]) -> Iterator[JSONLEvent]:
+        """Iterate over JSONL events from a byte iterator.
+
+        :param iter_bytes: An iterator of byte chunks.
+        :type iter_bytes: Iterator[bytes]
+        :rtype: Iterator[~azure.core.streaming.JSONLEvent]
+        :return: An iterator of JSONL events.
+        """
+
+        yield from (JSONLEvent(data=line) for line in iter_lines(iter_bytes))
+
+
+class AsyncJSONLDecoder:
+    """Asynchronous decoder for JSON Lines (JSONL) format. https://jsonlines.org/"""
+
+    # pylint: disable=invalid-overridden-method
+    async def aiter_events(self, iter_bytes: AsyncIterator[bytes]) -> AsyncIterator[JSONLEvent]:
+        """Asynchronously iterate over JSONL events from a byte iterator.
+
+        :param iter_bytes: An asynchronous iterator of byte chunks.
+        :type iter_bytes: AsyncIterator[bytes]
+        :rtype: AsyncIterator[~azure.core.streaming.JSONLEvent]
+        :return: An asynchronous iterator of JSONL events.
+        """
+
+        async with aclosing(aiter_lines(iter_bytes)) as lines:
+            async for line in lines:
+                yield JSONLEvent(data=line)