@typespec/http-client-python 0.37.2 → 0.38.0

package/generator/build/lib/pygen/codegen/templates/streaming_base.py.jinja2

@@ -33,7 +33,8 @@ from typing import (
 
 from typing_extensions import Self
 
-from {{ code_model.core_library }}.rest import AsyncHttpResponse, HttpResponse
+from {{ code_model.core_library }}.exceptions import StreamClosedError, StreamConsumedError
+from {{ code_model.core_library }}.rest import AsyncHttpResponse, HttpRequest, HttpResponse
 
 DecodedType = TypeVar("DecodedType")
 ReturnType_co = TypeVar("ReturnType_co", covariant=True)
@@ -435,13 +436,56 @@ async def _aiter_sse_lines(iter_bytes: AsyncIterator[bytes]) -> AsyncGenerator[s
         yield line
 
 
+def _read_sse_response(response: HttpResponse) -> None:
+    """Read an SSE response body, ignoring responses that are already consumed or closed.
+
+    :param response: The response to read.
+    :type response: HttpResponse
+    """
+    try:
+        response.read()
+    except (StreamConsumedError, StreamClosedError):
+        pass
+
+
+async def _read_sse_response_async(response: AsyncHttpResponse) -> None:
+    """Read an SSE response body, ignoring responses that are already consumed or closed.
+
+    :param response: The response to read.
+    :type response: AsyncHttpResponse
+    """
+    try:
+        await response.read()
+    except (StreamConsumedError, StreamClosedError):
+        pass
+
+
+def _update_sse_request_headers(request: HttpRequest, last_event_id: Optional[str]) -> None:
+    """Update the last event ID header for an SSE reconnection request.
+
+    :param request: The reconnection request to update.
+    :type request: HttpRequest
+    :param last_event_id: The current value of the SSE last event ID buffer.
+    :type last_event_id: str or None
+    """
+    if last_event_id:
+        request.headers["Last-Event-ID"] = last_event_id
+    else:
+        request.headers.pop("Last-Event-ID", None)
+
+
 class _SSEEventBuilder:
-    """Accumulates SSE field lines and builds :class:`ServerSentEvent` instances."""
+    """Accumulates SSE field lines and builds :class:`ServerSentEvent` instances.
 
-    def __init__(self) -> None:
+    :param last_event_id: The initial value of the SSE last event ID buffer.
+    :type last_event_id: str
+    """
+
+    def __init__(self, last_event_id: str = "") -> None:
         self._data: List[str] = []
         self._event_type = ""
-        self._last_id = ""
+        self._last_id = last_event_id
+        self._pending_last_id = last_event_id
         self._retry: Optional[int] = None
 
     def add_line(self, line: str) -> Optional[ServerSentEvent]:
@@ -469,7 +513,7 @@ class _SSEEventBuilder:
             self._data.append(value)
         elif field == "id":
             if "\x00" not in value:
-                self._last_id = value
+                self._pending_last_id = value
         elif field == "retry":
             if value.isascii() and value.isdigit():
                 try:
@@ -482,6 +526,7 @@ class _SSEEventBuilder:
         return None
 
     def _dispatch(self) -> Optional[ServerSentEvent]:
+        self._last_id = self._pending_last_id
         if not self._data:
             # No data accumulated: reset and dispatch nothing.
             self._event_type = ""
@@ -496,13 +541,55 @@ class _SSEEventBuilder:
         self._event_type = ""
         return event
 
+    @property
+    def last_event_id(self) -> str:
+        """Return the current value of the last event ID buffer.
+
+        :return: The current value of the last event ID buffer.
+        :rtype: str
+        """
+        return self._last_id
+
+    @property
+    def retry(self) -> Optional[int]:
+        """Return the most recently received valid reconnection time.
+
+        :return: The reconnection time in milliseconds, if one was provided.
+        :rtype: int or None
+        """
+        return self._retry
+
 
 class SSEDecoder:
     """Decoder for Server-Sent Events (SSE).
 
     https://html.spec.whatwg.org/multipage/server-sent-events.html
+
+    :keyword last_event_id: The initial value of the SSE last event ID buffer.
+    :paramtype last_event_id: str or None
     """
 
+    def __init__(self, *, last_event_id: Optional[str] = None) -> None:
+        self._builder = _SSEEventBuilder(last_event_id or "")
+
+    @property
+    def last_event_id(self) -> str:
+        """Return the current value of the last event ID buffer.
+
+        :return: The current value of the last event ID buffer.
+        :rtype: str
+        """
+        return self._builder.last_event_id
+
+    @property
+    def retry(self) -> Optional[int]:
+        """Return the most recently received valid reconnection time.
+
+        :return: The reconnection time in milliseconds, if one was provided.
+        :rtype: int or None
+        """
+        return self._builder.retry
+
     def iter_events(self, iter_bytes: Iterator[bytes]) -> Iterator[ServerSentEvent]:
         """Iterate over SSE events from a byte iterator.
 
@@ -511,9 +598,8 @@ class SSEDecoder:
         :rtype: Iterator[ServerSentEvent]
         :return: An iterator of server-sent events.
         """
-        builder = _SSEEventBuilder()
         for line in _iter_sse_lines(iter_bytes):
-            event = builder.add_line(line)
+            event = self._builder.add_line(line)
             if event is not None:
                 yield event
 
@@ -522,8 +608,32 @@ class AsyncSSEDecoder:
     """Asynchronous decoder for Server-Sent Events (SSE).
 
     https://html.spec.whatwg.org/multipage/server-sent-events.html
+
+    :keyword last_event_id: The initial value of the SSE last event ID buffer.
+    :paramtype last_event_id: str or None
     """
 
+    def __init__(self, *, last_event_id: Optional[str] = None) -> None:
+        self._builder = _SSEEventBuilder(last_event_id or "")
+
+    @property
+    def last_event_id(self) -> str:
+        """Return the current value of the last event ID buffer.
+
+        :return: The current value of the last event ID buffer.
+        :rtype: str
+        """
+        return self._builder.last_event_id
+
+    @property
+    def retry(self) -> Optional[int]:
+        """Return the most recently received valid reconnection time.
+
+        :return: The reconnection time in milliseconds, if one was provided.
+        :rtype: int or None
+        """
+        return self._builder.retry
+
     # pylint: disable=invalid-overridden-method
     async def aiter_events(self, iter_bytes: AsyncIterator[bytes]) -> AsyncIterator[ServerSentEvent]:
         """Asynchronously iterate over SSE events from a byte iterator.
@@ -533,10 +643,9 @@ class AsyncSSEDecoder:
         :rtype: AsyncIterator[ServerSentEvent]
         :return: An asynchronous iterator of server-sent events.
         """
-        builder = _SSEEventBuilder()
         async with aclosing(_aiter_sse_lines(iter_bytes)) as lines:
             async for line in lines:
-                event = builder.add_line(line)
+                event = self._builder.add_line(line)
                 if event is not None:
                     yield event
 
@@ -563,6 +672,14 @@ class Stream(Iterator[ReturnType_co]):
     :keyword terminal_event_predicate: Optional predicate that identifies a payload-bearing
         terminal event. The event is yielded before iteration stops.
     :paramtype terminal_event_predicate: Callable[[Any], bool] or None
+    :keyword last_event_id: Optional event ID from which this stream is resuming.
+    :paramtype last_event_id: str or None
+    :keyword reconnect_callback: Optional callback used to re-establish the underlying response
+        when the stream disconnects before a terminal event is seen. Reconnection uses a
+        three-second delay unless the server provides an SSE ``retry`` value. The callback is
+        called with the last-seen SSE event ID (or ``None``) and the reconnect delay in seconds,
+        and must return a new response.
+    :paramtype reconnect_callback: Callable[[str or None, float], ~{{ code_model.core_library }}.rest.HttpResponse] or None
     """
 
     def __init__(
@@ -574,17 +691,26 @@ class Stream(Iterator[ReturnType_co]):
         terminal_event: Optional[str] = None,
         terminal_event_names: Optional[Sequence[str]] = None,
         terminal_event_predicate: Optional[Callable[[DecodedType], bool]] = None,
+        last_event_id: Optional[str] = None,
+        reconnect_callback: Optional[Callable[[Optional[str], float], HttpResponse]] = None,
     ) -> None:
         self._response = response
         content_type = response.headers.get("Content-Type", "").split(";", 1)[0].strip().lower()
         self._decoder: StreamDecoder[Any] = (
-            decoder if decoder is not None else (SSEDecoder() if content_type == "text/event-stream" else JSONLDecoder())
+            decoder
+            if decoder is not None
+            else (
+                SSEDecoder(last_event_id=last_event_id)
+                if content_type == "text/event-stream"
+                else JSONLDecoder()
+            )
         )
         self._deserialization_callback = deserialization_callback
         self._terminal_event = terminal_event
         self._terminal_event_names = frozenset(terminal_event_names or ())
         self._terminal_event_predicate = terminal_event_predicate
-        self._last_event_id: Optional[str] = None
+        self._reconnect_callback = reconnect_callback
+        self._last_event_id = last_event_id
         self._retry: Optional[int] = None
         self._iterator = self._iter_results()
 
@@ -605,26 +731,46 @@ class Stream(Iterator[ReturnType_co]):
         return self
 
     def _iter_results(self) -> Generator[ReturnType_co, None, None]:
+        terminal = False
         try:
-            for event in self._decoder.iter_events(self._response.iter_bytes()):
-                event_id = getattr(event, "id", None)
-                if event_id is not None:
-                    self._last_event_id = event_id
-                event_retry = getattr(event, "retry", None)
-                if event_retry is not None:
-                    self._retry = event_retry
-                if self._terminal_event is not None and getattr(event, "data", None) == self._terminal_event:
+            while True:
+                if getattr(self._response, "status_code", None) == 204:
                     break
-                result = self._deserialization_callback(self._response, event)
-                yield result
-                if (
-                    self._terminal_event_names
-                    and getattr(event, "event", None) in self._terminal_event_names
-                ) or (
-                    self._terminal_event_predicate is not None
-                    and self._terminal_event_predicate(event)
-                ):
+                for event in self._decoder.iter_events(self._response.iter_bytes()):
+                    event_id = getattr(event, "id", None)
+                    if event_id is not None:
+                        self._last_event_id = event_id
+                    event_retry = getattr(event, "retry", None)
+                    if event_retry is not None:
+                        self._retry = event_retry
+                    if self._terminal_event is not None and getattr(event, "data", None) == self._terminal_event:
+                        terminal = True
+                        break
+                    result = self._deserialization_callback(self._response, cast(Any, event))
+                    yield result
+                    if (
+                        self._terminal_event_names
+                        and getattr(event, "event", None) in self._terminal_event_names
+                    ) or (
+                        self._terminal_event_predicate is not None
+                        and self._terminal_event_predicate(cast(Any, event))
+                    ):
+                        terminal = True
+                        break
+                if terminal or self._reconnect_callback is None:
                     break
+                if isinstance(self._decoder, SSEDecoder):
+                    self._last_event_id = self._decoder.last_event_id
+                    if self._decoder.retry is not None:
+                        self._retry = self._decoder.retry
+                self._response.close()
+                reconnect_delay = (self._retry if self._retry is not None else 3000) / 1000
+                self._response = self._reconnect_callback(self._last_event_id, reconnect_delay)
+                self._decoder = (
+                    SSEDecoder(last_event_id=self._last_event_id)
+                    if isinstance(self._decoder, SSEDecoder)
+                    else type(self._decoder)()
+                )
         finally:
             self._response.close()
 
@@ -668,6 +814,14 @@ class AsyncStream(AsyncIterator[ReturnType_co]):
     :keyword terminal_event_predicate: Optional predicate that identifies a payload-bearing
         terminal event. The event is yielded before iteration stops.
     :paramtype terminal_event_predicate: Callable[[Any], bool] or None
+    :keyword last_event_id: Optional event ID from which this stream is resuming.
+    :paramtype last_event_id: str or None
+    :keyword reconnect_callback: Optional callback used to re-establish the underlying response
+        when the stream disconnects before a terminal event is seen. Reconnection uses a
+        three-second delay unless the server provides an SSE ``retry`` value. The callback is
+        called with the last-seen SSE event ID (or ``None``) and the reconnect delay in seconds,
+        and must resolve to a new response.
+    :paramtype reconnect_callback: Callable[[str or None, float], Any] or None
     """
 
     def __init__(
@@ -679,19 +833,26 @@ class AsyncStream(AsyncIterator[ReturnType_co]):
         terminal_event: Optional[str] = None,
         terminal_event_names: Optional[Sequence[str]] = None,
         terminal_event_predicate: Optional[Callable[[DecodedType], bool]] = None,
+        last_event_id: Optional[str] = None,
+        reconnect_callback: Optional[Callable[[Optional[str], float], Any]] = None,
     ) -> None:
         self._response = response
         content_type = response.headers.get("Content-Type", "").split(";", 1)[0].strip().lower()
         self._decoder: AsyncStreamDecoder[Any] = (
             decoder
             if decoder is not None
-            else (AsyncSSEDecoder() if content_type == "text/event-stream" else AsyncJSONLDecoder())
+            else (
+                AsyncSSEDecoder(last_event_id=last_event_id)
+                if content_type == "text/event-stream"
+                else AsyncJSONLDecoder()
+            )
         )
         self._deserialization_callback = deserialization_callback
         self._terminal_event = terminal_event
         self._terminal_event_names = frozenset(terminal_event_names or ())
         self._terminal_event_predicate = terminal_event_predicate
-        self._last_event_id: Optional[str] = None
+        self._reconnect_callback = reconnect_callback
+        self._last_event_id = last_event_id
         self._retry: Optional[int] = None
         self._iterator = self._iter_results()
 
@@ -712,34 +873,54 @@ class AsyncStream(AsyncIterator[ReturnType_co]):
         return self
 
     async def _iter_results(self) -> AsyncGenerator[ReturnType_co, None]:
-        events = self._decoder.aiter_events(self._response.iter_bytes())
+        terminal = False
         try:
-            async for event in events:
-                event_id = getattr(event, "id", None)
-                if event_id is not None:
-                    self._last_event_id = event_id
-                event_retry = getattr(event, "retry", None)
-                if event_retry is not None:
-                    self._retry = event_retry
-                if self._terminal_event is not None and getattr(event, "data", None) == self._terminal_event:
+            while True:
+                if getattr(self._response, "status_code", None) == 204:
                     break
-                result = self._deserialization_callback(self._response, event)
-                yield result
-                if (
-                    self._terminal_event_names
-                    and getattr(event, "event", None) in self._terminal_event_names
-                ) or (
-                    self._terminal_event_predicate is not None
-                    and self._terminal_event_predicate(event)
-                ):
+                events = self._decoder.aiter_events(self._response.iter_bytes())
+                try:
+                    async for event in events:
+                        event_id = getattr(event, "id", None)
+                        if event_id is not None:
+                            self._last_event_id = event_id
+                        event_retry = getattr(event, "retry", None)
+                        if event_retry is not None:
+                            self._retry = event_retry
+                        if self._terminal_event is not None and getattr(event, "data", None) == self._terminal_event:
+                            terminal = True
+                            break
+                        result = self._deserialization_callback(self._response, cast(Any, event))
+                        yield result
+                        if (
+                            self._terminal_event_names
+                            and getattr(event, "event", None) in self._terminal_event_names
+                        ) or (
+                            self._terminal_event_predicate is not None
+                            and self._terminal_event_predicate(cast(Any, event))
+                        ):
+                            terminal = True
+                            break
+                finally:
+                    aclose = getattr(events, "aclose", None)
+                    if aclose is not None:
+                        await aclose()
+                if terminal or self._reconnect_callback is None:
                     break
-        finally:
-            try:
-                aclose = getattr(events, "aclose", None)
-                if aclose is not None:
-                    await aclose()
-            finally:
+                if isinstance(self._decoder, AsyncSSEDecoder):
+                    self._last_event_id = self._decoder.last_event_id
+                    if self._decoder.retry is not None:
+                        self._retry = self._decoder.retry
                 await self._response.close()
+                reconnect_delay = (self._retry if self._retry is not None else 3000) / 1000
+                self._response = await self._reconnect_callback(self._last_event_id, reconnect_delay)
+                self._decoder = (
+                    AsyncSSEDecoder(last_event_id=self._last_event_id)
+                    if isinstance(self._decoder, AsyncSSEDecoder)
+                    else type(self._decoder)()
+                )
+        finally:
+            await self._response.close()
 
     async def __aexit__(
         self,

package/generator/build/lib/pygen/codegen/templates/unions.py.jinja2

@@ -4,6 +4,6 @@
 {% endif %}
 
 {{ imports }}
-{% for nu in code_model.named_unions %}
-{{nu.name}} = {{nu.type_definition()}}
+{% for nu in named_unions %}
+{{nu.name}}: TypeAlias = {{nu.type_definition()}}
 {% endfor %}

package/generator/build/lib/pygen/preprocess/__init__.py

@@ -54,13 +54,14 @@ def add_overload(yaml_data: dict[str, Any], body_type: dict[str, Any], for_flatt
     if yaml_data.get("initialOperation"):
         overload["initialOperation"] = yaml_data["initialOperation"]
 
+    # Reattach shared type objects before filtering parameters so positional
+    # alignment with the original operation is preserved.
+    for overload_p, original_p in zip(overload["parameters"], yaml_data["parameters"]):
+        overload_p["type"] = original_p["type"]
     if for_flatten_params:
         overload["bodyParameter"]["flattened"] = True
     else:
         overload["parameters"] = [p for p in overload["parameters"] if not p.get("inFlattenedBody")]
-    # for yaml sync, we need to make sure all of the responses, parameters, and exceptions' types have the same yaml id
-    for overload_p, original_p in zip(overload["parameters"], yaml_data["parameters"]):
-        overload_p["type"] = original_p["type"]
     update_overload_section(overload, yaml_data, "responses")
     update_overload_section(overload, yaml_data, "exceptions")
 
@@ -95,12 +96,14 @@ def add_overloads_for_body_param(yaml_data: dict[str, Any], skip_single_body_jso
     raw-JSON overload is kept, matching pre-TypedDict behavior).
     """
     body_parameter = yaml_data["bodyParameter"]
+    body_types = body_parameter["type"].get("types", [])
     if not (
         body_parameter["type"]["type"] == "combined"
-        and len(yaml_data["bodyParameter"]["type"]["types"]) > len(yaml_data["overloads"])
+        and len(body_types) > 1
+        and len(body_types) > len(yaml_data["overloads"])
     ):
         return
-    for body_type in body_parameter["type"]["types"]:
+    for body_type in body_types:
         if any(o for o in yaml_data["overloads"] if id(o["bodyParameter"]["type"]) == id(body_type)):
             continue
         if body_type.get("type") == "model" and body_type.get("base") == "json":
@@ -262,9 +265,6 @@ def _process_operation_etag_headers(
     if_match_candidates: list[dict[str, Any]] = []
     if_none_match_candidates: list[dict[str, Any]] = []
     for p in operation["parameters"]:
-        wire_name_lower = get_wire_name_lower(p)
-        if p["location"] == "header" and wire_name_lower == "client-request-id":
-            client["requestIdHeaderName"] = wire_name_lower
         if version_tolerant and p["location"] == "header":
             role = _get_etag_role(p)
             if role == "ifMatch":
@@ -273,7 +273,6 @@ def _process_operation_etag_headers(
                 if_none_match_candidates.append(p)
 
     property_if_match, property_if_none_match = _resolve_etag_pair(if_match_candidates, if_none_match_candidates)
-
     if property_if_match and property_if_none_match:
         etag_params = {id(property_if_match), id(property_if_none_match)}
         operation["parameters"] = [item for item in operation["parameters"] if id(item) not in etag_params] + [
@@ -284,6 +283,21 @@ def _process_operation_etag_headers(
         client["hasEtag"] = True
 
 
+def _process_operation_group_etag_headers(
+    operation_groups: list[dict[str, Any]],
+    client: dict[str, Any],
+    version_tolerant: bool,
+) -> None:
+    for operation_group in operation_groups:
+        for operation in operation_group.get("operations", []):
+            _process_operation_etag_headers(operation, client, version_tolerant)
+        _process_operation_group_etag_headers(
+            operation_group.get("operationGroups", []),
+            client,
+            version_tolerant,
+        )
+
+
 def headers_convert(yaml_data: dict[str, Any], replace_data: Any) -> None:
     if isinstance(replace_data, dict):
         for k, v in replace_data.items():
@@ -605,9 +619,16 @@ class PreProcessPlugin(YamlUpdatePlugin):
         if prop_name.endswith("Client"):
             prop_name = prop_name[: len(prop_name) - len("Client")]
         yaml_data["builderPadName"] = to_snake_case(prop_name)
-        for og in yaml_data.get("operationGroups", []):
-            for o in og["operations"]:
-                _process_operation_etag_headers(o, yaml_data, self.version_tolerant)
+        for operation_group in yaml_data.get("operationGroups", []):
+            for operation in operation_group.get("operations", []):
+                for parameter in operation["parameters"]:
+                    if parameter["location"] == "header" and get_wire_name_lower(parameter) == "client-request-id":
+                        yaml_data["requestIdHeaderName"] = "client-request-id"
+        _process_operation_group_etag_headers(
+            yaml_data.get("operationGroups", []),
+            yaml_data,
+            self.version_tolerant,
+        )
 
         # add client signature cloud_setting for arm
         if self.azure_arm and yaml_data["parameters"]:

package/generator/pygen/codegen/models/response.py

@@ -267,6 +267,11 @@ class Response(BaseModel):
             file_import.add_submodule_import(relative_path, stream_class, ImportType.LOCAL)
             if self.streaming_kind == "sse":
                 file_import.add_import("json", ImportType.STDLIB)
+                read_response_helper = (
+                    "_read_sse_response_async" if kwargs.get("async_mode", False) else "_read_sse_response"
+                )
+                file_import.add_submodule_import(relative_path, read_response_helper, ImportType.LOCAL)
+                file_import.add_submodule_import(relative_path, "_update_sse_request_headers", ImportType.LOCAL)
                 for event in self.streaming_events:
                     file_import.merge(event.payload_type.imports(**kwargs))
         return file_import