truss 0.10.0rc1__py3-none-any.whl → 0.60.0__py3-none-any.whl

truss_chains/remote_chainlet/utils.py

@@ -0,0 +1,332 @@
+import asyncio
+import builtins
+import contextlib
+import contextvars
+import json
+import logging
+import sys
+import textwrap
+import threading
+import traceback
+from typing import Dict, Iterator, Mapping, NoReturn, Optional, Type, TypeVar
+
+import aiohttp
+import fastapi
+import httpx
+import pydantic
+import starlette.requests
+from truss.templates.shared import dynamic_config_resolver
+
+from truss_chains import definitions
+
+T = TypeVar("T")
+
+
+def populate_chainlet_service_predict_urls(
+    chainlet_to_service: Mapping[str, definitions.ServiceDescriptor],
+) -> Mapping[str, definitions.DeployedServiceDescriptor]:
+    chainlet_to_deployed_service: Dict[str, definitions.DeployedServiceDescriptor] = {}
+    # If there are no dependencies of this chainlet, no need to derive dynamic URLs
+    if len(chainlet_to_service) == 0:
+        return chainlet_to_deployed_service
+
+    dynamic_chainlet_config_str = dynamic_config_resolver.get_dynamic_config_value_sync(
+        definitions.DYNAMIC_CHAINLET_CONFIG_KEY
+    )
+
+    if not dynamic_chainlet_config_str:
+        raise definitions.MissingDependencyError(
+            f"No '{definitions.DYNAMIC_CHAINLET_CONFIG_KEY}' "
+            "found. Cannot override Chainlet configs."
+        )
+
+    dynamic_chainlet_config = json.loads(dynamic_chainlet_config_str)
+
+    for chainlet_name, service_descriptor in chainlet_to_service.items():
+        display_name = service_descriptor.display_name
+
+        # NOTE: The Chainlet `display_name` in the Truss CLI
+        # corresponds to Chainlet `name` in the backend. As
+        # the dynamic Chainlet config is keyed on the backend
+        # Chainlet name, we have to look up config values by
+        # using the `display_name` in the service descriptor.
+        if display_name not in dynamic_chainlet_config:
+            raise definitions.MissingDependencyError(
+                f"Chainlet '{display_name}' not found in "
+                f"'{definitions.DYNAMIC_CHAINLET_CONFIG_KEY}'. "
+                f"Dynamic Chainlet config keys: {list(dynamic_chainlet_config)}."
+            )
+
+        chainlet_to_deployed_service[chainlet_name] = (
+            definitions.DeployedServiceDescriptor(
+                display_name=display_name,
+                name=service_descriptor.name,
+                options=service_descriptor.options,
+                predict_url=dynamic_chainlet_config[display_name]["predict_url"],
+            )
+        )
+
+    return chainlet_to_deployed_service
+
+
+class AsyncSafeCounter:
+    def __init__(self, initial: int = 0) -> None:
+        self._counter = initial
+        self._lock = asyncio.Lock()
+
+    async def increment(self) -> int:
+        async with self._lock:
+            self._counter += 1
+            return self._counter
+
+    async def decrement(self) -> int:
+        async with self._lock:
+            self._counter -= 1
+            return self._counter
+
+    async def __aenter__(self) -> int:
+        return await self.increment()
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        await self.decrement()
+
+
+class ThreadSafeCounter:
+    def __init__(self, initial: int = 0) -> None:
+        self._counter = initial
+        self._lock = threading.Lock()
+
+    def increment(self) -> int:
+        with self._lock:
+            self._counter += 1
+            return self._counter
+
+    def decrement(self) -> int:
+        with self._lock:
+            self._counter -= 1
+            return self._counter
+
+    def __enter__(self) -> int:
+        return self.increment()
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        self.decrement()
+
+
+_trace_parent_context: contextvars.ContextVar[str] = contextvars.ContextVar(
+    "trace_parent"
+)
+
+
+@contextlib.contextmanager
+def _trace_parent(request: starlette.requests.Request) -> Iterator[None]:
+    token = _trace_parent_context.set(
+        request.headers.get(definitions.OTEL_TRACE_PARENT_HEADER_KEY, "")
+    )
+    try:
+        yield
+    finally:
+        _trace_parent_context.reset(token)
+
+
+@contextlib.contextmanager
+def trace_parent_raw(trace_parent: str) -> Iterator[None]:
+    token = _trace_parent_context.set(trace_parent)
+    try:
+        yield
+    finally:
+        _trace_parent_context.reset(token)
+
+
+def get_trace_parent() -> Optional[str]:
+    return _trace_parent_context.get()
+
+
+def pydantic_set_field_dict(obj: pydantic.BaseModel) -> dict[str, pydantic.BaseModel]:
+    """Like `BaseModel.model_dump(exclude_unset=True), but only top-level.
+
+    This is used to get kwargs for invoking a function, while dropping fields for which
+    there is no value explicitly set in the pydantic model. A field is considered unset
+    if the key was not present in the incoming JSON request (from which the model was
+    parsed/initialized) and the pydantic model has a default value, such as `None`.
+
+    By dropping these unset fields, the default values from the function definition
+    will be used instead. This behavior ensures correct handling of arguments where
+    the function has a default, such as in the case of `run_remote`. If the model has
+    an optional field defaulting to `None`, this approach differentiates between
+    the user explicitly passing a value of `None` and the field being unset in the
+    request.
+
+    """
+    return {name: getattr(obj, name) for name in obj.model_fields_set}
+
+
+# Error Propagation Utils. #############################################################
+
+
+def _handle_exception(exception: Exception) -> NoReturn:
+    """Raises `HTTPException` with `RemoteErrorDetail`."""
+    if hasattr(exception, "__module__"):
+        exception_module_name = exception.__module__
+    else:
+        exception_module_name = None
+
+    error_stack = traceback.extract_tb(exception.__traceback__)
+    # Filter everything before (model.py) and after (stubs, error handling) so that only
+    # user-defined code remains. See test_e2e.py::test_chain for expected results.
+    model_predict_index = 0
+    first_stub_index = len(error_stack)
+    for i, frame in enumerate(error_stack):
+        if frame.filename.endswith("model/model.py") and frame.name == "predict":
+            model_predict_index = i + 1
+        if frame.filename.endswith("remote_chainlet/stub.py") and frame.name.startswith(
+            "predict"  # predict sycnc|async|stream.
+        ):
+            first_stub_index = i - 1
+            break
+
+    final_tb = error_stack[model_predict_index:first_stub_index]
+    stack = [definitions.StackFrame.from_frame_summary(frame) for frame in final_tb]
+    error = definitions.RemoteErrorDetail(
+        exception_cls_name=exception.__class__.__name__,
+        exception_module_name=exception_module_name,
+        exception_message=str(exception),
+        user_stack_trace=list(stack),
+    )
+    raise fastapi.HTTPException(
+        status_code=500, detail=error.model_dump()
+    ) from exception
+
+
+@contextlib.contextmanager
+def _exception_to_http_error() -> Iterator[None]:
+    try:
+        yield
+    except Exception as e:
+        _handle_exception(e)
+
+
+def _resolve_exception_class(error: definitions.RemoteErrorDetail) -> Type[Exception]:
+    """Tries to find the exception class in builtins or imported libs,
+    falls back to `definitions.GenericRemoteError` if not found."""
+    exception_cls = None
+    if error.exception_module_name is None:
+        exception_cls = getattr(builtins, error.exception_cls_name, None)
+    else:
+        if mod := sys.modules.get(error.exception_module_name):
+            exception_cls = getattr(mod, error.exception_cls_name, None)
+
+    if exception_cls is None:
+        logging.warning(
+            f"Could not resolve exception with name `{error.exception_cls_name}` "
+            f"and module `{error.exception_module_name}` - fall back to "
+            f"`{definitions.GenericRemoteException.__name__}`."
+        )
+        exception_cls = definitions.GenericRemoteException
+
+    if issubclass(exception_cls, pydantic.ValidationError):
+        # Cannot re-raise naively.
+        # https://github.com/pydantic/pydantic/issues/6734.
+        exception_cls = definitions.GenericRemoteException
+
+    return exception_cls
+
+
+def _handle_response_error(response_json: dict, base_msg: str):
+    try:
+        error_json = response_json["error"]
+    except KeyError as e:
+        logging.error(f"response_json: {response_json}")
+        raise ValueError(
+            f"{base_msg}. Could not get `error` field from JSON response."
+        ) from e
+
+    try:
+        error = definitions.RemoteErrorDetail.model_validate(error_json)
+    except pydantic.ValidationError as e:
+        if isinstance(error_json, str):
+            msg = f"{base_msg}: '{error_json}'"
+            raise definitions.GenericRemoteException(msg) from None
+        raise ValueError(
+            f"{base_msg}: Could not parse chainlet error. Error details are expected "
+            "to be either a plain string (old truss models) or a serialized "
+            f"`{definitions.RemoteErrorDetail.__name__}`, got:\n{repr(error_json)}"
+        ) from e
+
+    exception_cls = _resolve_exception_class(error)
+    error_format = textwrap.indent(error.format(), "│   ")
+    *lines, last_line = error_format.splitlines()
+    last_line = f"╰{last_line[1:]}" if last_line.startswith("│") else last_line
+    error_format = "\n".join(lines + [last_line])
+    msg = (
+        f"(showing chained remote errors, root error at the bottom)\n"
+        f"├─ {base_msg}\n"
+        f"{error_format}"
+    )
+    raise exception_cls(msg)
+
+
+def _make_base_error_message(remote_name: str, http_status: int) -> str:
+    return (
+        f"Error calling dependency Chainlet `{remote_name}`, "
+        f"HTTP status={http_status}, trace ID=`{get_trace_parent()}`."
+    )
+
+
+def response_raise_errors(response: httpx.Response, remote_name: str) -> None:
+    """In case of error, raise it.
+
+    If the response error contains `RemoteErrorDetail`, it tries to re-raise
+    the same exception that was raised remotely and falls back to
+    `GenericRemoteException` if the exception class could not be resolved.
+
+    Exception messages are chained to trace back to the root cause, i.e. the first
+    Chainlet that raised an exception. E.g. the message might look like this:
+
+    ```
+    Chainlet-Traceback (most recent call last):
+      File "/packages/itest_chain.py", line 132, in run_remote
+        value = self._accumulate_parts(text_parts.parts)
+      File "/packages/itest_chain.py", line 144, in _accumulate_parts
+        value += self._text_to_num.run_remote(part)
+    ValueError: (showing chained remote errors, root error at the bottom)
+    ├─ Error in dependency Chainlet `TextToNum` (HTTP status 500):
+    │   Chainlet-Traceback (most recent call last):
+    │     File "/packages/itest_chain.py", line 87, in run_remote
+    │       generated_text = self._replicator.run_remote(data)
+    │   ValueError: (showing chained remote errors, root error at the bottom)
+    │   ├─ Error in dependency Chainlet `TextReplicator` (HTTP status 500):
+    │   │   Chainlet-Traceback (most recent call last):
+    │   │     File "/packages/itest_chain.py", line 52, in run_remote
+    │   │       validate_data(data)
+    │   │     File "/packages/itest_chain.py", line 36, in validate_data
+    │   │       raise ValueError(f"This input is too long: {len(data)}.")
+    ╰   ╰   ValueError: This input is too long: 100.
+    ```
+    """
+    if response.is_error:
+        base_msg = _make_base_error_message(remote_name, response.status_code)
+        try:
+            response_json = response.json()
+        except Exception as e:
+            raise ValueError(base_msg) from e
+        _handle_response_error(response_json, base_msg)
+
+
+async def async_response_raise_errors(
+    response: aiohttp.ClientResponse, remote_name: str
+) -> None:
+    """Async version of `async_response_raise_errors`."""
+    if response.status >= 400:
+        base_msg = _make_base_error_message(remote_name, response.status)
+        try:
+            response_json = await response.json()
+        except Exception as e:
+            raise ValueError(base_msg) from e
+        _handle_response_error(response_json, base_msg)
+
+
+@contextlib.contextmanager
+def predict_context(request: starlette.requests.Request) -> Iterator[None]:
+    with _trace_parent(request), _exception_to_http_error():
+        yield

truss_chains/streaming.py

@@ -0,0 +1,378 @@
+import asyncio
+import dataclasses
+import enum
+import struct
+import sys
+from collections.abc import AsyncIterator
+from typing import Generic, Optional, Protocol, Type, TypeVar, overload
+
+import pydantic
+
+_TAG_SIZE = 5  # uint8 + uint32.
+
+_T = TypeVar("_T")
+
+if sys.version_info < (3, 10):
+
+    async def anext(iterable: AsyncIterator[_T]) -> _T:
+        return await iterable.__anext__()
+
+
+# Note on the (verbose) typing in this module: we want exact typing of the reader and
+# writer helpers, while also allowing flexibility to users to leave out header/footer
+# if not needed.
+# Putting both a constraint on the header/footer types to be pydantic
+# models, but also letting them be optional is not well-supported by typing tools,
+# (missing feature is using type variables a constraints on other type variables).
+#
+# A functional, yet verbose workaround that gives correct variadic type inference,
+# is using intermediate type variables `HeaderT` <-> `HeaderTT` and in conjunction with
+# mapping out all usage combinations with overloads (the overloads essentially allow
+# "conditional" binding of type vars). These overloads also allow to use granular
+# reader/writer sub-classes conditionally, that have the read/write methods only for the
+# data types configured, and implemented DRY with mixin classes.
+ItemT = TypeVar("ItemT", bound=pydantic.BaseModel)
+HeaderT = TypeVar("HeaderT", bound=pydantic.BaseModel)
+FooterT = TypeVar("FooterT", bound=pydantic.BaseModel)
+
+# Since header/footer could also be `None`, we need an extra type variable that
+# can assume either `Type[HeaderT]` or `None` - `Type[None]` causes issues.
+HeaderTT = TypeVar("HeaderTT")
+FooterTT = TypeVar("FooterTT")
+
+
+@dataclasses.dataclass
+class StreamTypes(Generic[ItemT, HeaderTT, FooterTT]):
+    item_type: Type[ItemT]
+    header_type: HeaderTT  # Is either `Type[HeaderT]` or `None`.
+    footer_type: FooterTT  # Is either `Type[FooterT]` or `None`.
+
+
+@overload
+def stream_types(
+    item_type: Type[ItemT], *, header_type: Type[HeaderT], footer_type: Type[FooterT]
+) -> StreamTypes[ItemT, HeaderT, FooterT]: ...
+
+
+@overload
+def stream_types(
+    item_type: Type[ItemT], *, header_type: Type[HeaderT]
+) -> StreamTypes[ItemT, HeaderT, None]: ...
+
+
+@overload
+def stream_types(
+    item_type: Type[ItemT], *, footer_type: Type[FooterT]
+) -> StreamTypes[ItemT, None, FooterT]: ...
+
+
+@overload
+def stream_types(item_type: Type[ItemT]) -> StreamTypes[ItemT, None, None]: ...
+
+
+def stream_types(
+    item_type: Type[ItemT],
+    *,
+    header_type: Optional[Type[HeaderT]] = None,
+    footer_type: Optional[Type[FooterT]] = None,
+) -> StreamTypes:
+    """Creates a bundle of item type and potentially header/footer types,
+    each as pydantic model."""
+    # This indirection for creating `StreamTypes` is needed to get generic typing.
+    return StreamTypes(item_type, header_type, footer_type)
+
+
+# Reading ##############################################################################
+
+
+class _Delimiter(enum.IntEnum):
+    NOT_SET = enum.auto()
+    HEADER = enum.auto()
+    ITEM = enum.auto()
+    FOOTER = enum.auto()
+    END = enum.auto()
+
+
+class _Streamer(Generic[ItemT, HeaderTT, FooterTT]):
+    _stream_types: StreamTypes[ItemT, HeaderTT, FooterTT]
+
+    def __init__(self, types: StreamTypes[ItemT, HeaderTT, FooterTT]) -> None:
+        self._stream_types = types
+
+
+# Reading ##############################################################################
+
+
+class _ByteReader:
+    """Helper to provide `readexactly` API for an async bytes iterator."""
+
+    def __init__(self, source: AsyncIterator[bytes]) -> None:
+        self._source = source
+        self._buffer = bytearray()
+
+    async def readexactly(self, num_bytes: int) -> bytes:
+        while len(self._buffer) < num_bytes:
+            try:
+                chunk = await anext(self._source)
+            except StopAsyncIteration:
+                break
+            self._buffer.extend(chunk)
+
+        if len(self._buffer) < num_bytes:
+            if len(self._buffer) == 0:
+                raise EOFError()
+            raise asyncio.IncompleteReadError(self._buffer, num_bytes)
+
+        result = bytes(self._buffer[:num_bytes])
+        del self._buffer[:num_bytes]
+        return result
+
+
+class _StreamReaderProtocol(Protocol[ItemT, HeaderTT, FooterTT]):
+    _stream_types: StreamTypes[ItemT, HeaderTT, FooterTT]
+    _footer_data: Optional[bytes]
+
+    async def _read(self) -> tuple[_Delimiter, bytes]: ...
+
+
+class _StreamReader(_Streamer[ItemT, HeaderTT, FooterTT]):
+    _stream: _ByteReader
+    _footer_data: Optional[bytes]
+
+    def __init__(
+        self,
+        types: StreamTypes[ItemT, HeaderTT, FooterTT],
+        stream: AsyncIterator[bytes],
+    ) -> None:
+        super().__init__(types)
+        self._stream = _ByteReader(stream)
+        self._footer_data = None
+
+    @staticmethod
+    def _unpack_tag(tag: bytes) -> tuple[_Delimiter, int]:
+        enum_value, length = struct.unpack(">BI", tag)
+        return _Delimiter(enum_value), length
+
+    async def _read(self) -> tuple[_Delimiter, bytes]:
+        try:
+            tag = await self._stream.readexactly(_TAG_SIZE)
+        # It's ok to read nothing (end of stream), but unexpected to read partial.
+        except asyncio.IncompleteReadError:
+            raise
+        except EOFError:
+            return _Delimiter.END, b""
+
+        delimiter, length = self._unpack_tag(tag)
+        if not length:
+            return delimiter, b""
+        data_bytes = await self._stream.readexactly(length)
+        return delimiter, data_bytes
+
+    async def read_items(self) -> AsyncIterator[ItemT]:
+        delimiter, data_bytes = await self._read()
+        if delimiter == _Delimiter.HEADER:
+            raise ValueError(
+                "Called `read_items`, but there the stream contains header data, which "
+                "is not consumed. Call `read_header` first or remove sending a header."
+            )
+        if delimiter in (_Delimiter.FOOTER, _Delimiter.END):  # In case of 0 items.
+            self._footer_data = data_bytes
+            return
+
+        assert delimiter == _Delimiter.ITEM
+        while True:
+            yield self._stream_types.item_type.model_validate_json(data_bytes)
+            # We don't know if the next data is another item, footer or the end.
+            delimiter, data_bytes = await self._read()
+            if delimiter == _Delimiter.END:
+                return
+            if delimiter == _Delimiter.FOOTER:
+                self._footer_data = data_bytes
+                return
+
+
+class _HeaderReadMixin(_Streamer[ItemT, HeaderT, FooterTT]):
+    async def read_header(
+        self: _StreamReaderProtocol[ItemT, HeaderT, FooterTT],
+    ) -> HeaderT:
+        delimiter, data_bytes = await self._read()
+        if delimiter != _Delimiter.HEADER:
+            raise ValueError("Stream does not contain header.")
+        return self._stream_types.header_type.model_validate_json(data_bytes)
+
+
+class _FooterReadMixin(_Streamer[ItemT, HeaderTT, FooterT]):
+    _footer_data: Optional[bytes]
+
+    async def read_footer(
+        self: _StreamReaderProtocol[ItemT, HeaderTT, FooterT],
+    ) -> FooterT:
+        if self._footer_data is None:
+            delimiter, data_bytes = await self._read()
+            if delimiter != _Delimiter.FOOTER:
+                raise ValueError("Stream does not contain footer.")
+            self._footer_data = data_bytes
+
+        footer = self._stream_types.footer_type.model_validate_json(self._footer_data)
+        self._footer_data = None
+        return footer
+
+
+class StreamReaderWithHeader(
+    _StreamReader[ItemT, HeaderT, FooterTT], _HeaderReadMixin[ItemT, HeaderT, FooterTT]
+): ...
+
+
+class StreamReaderWithFooter(
+    _StreamReader[ItemT, HeaderTT, FooterT], _FooterReadMixin[ItemT, HeaderTT, FooterT]
+): ...
+
+
+class StreamReaderFull(
+    _StreamReader[ItemT, HeaderT, FooterT],
+    _HeaderReadMixin[ItemT, HeaderT, FooterT],
+    _FooterReadMixin[ItemT, HeaderT, FooterT],
+): ...
+
+
+@overload
+def stream_reader(
+    types: StreamTypes[ItemT, None, None], stream: AsyncIterator[bytes]
+) -> _StreamReader[ItemT, None, None]: ...
+
+
+@overload
+def stream_reader(
+    types: StreamTypes[ItemT, HeaderT, None], stream: AsyncIterator[bytes]
+) -> StreamReaderWithHeader[ItemT, HeaderT, None]: ...
+
+
+@overload
+def stream_reader(
+    types: StreamTypes[ItemT, None, FooterT], stream: AsyncIterator[bytes]
+) -> StreamReaderWithFooter[ItemT, None, FooterT]: ...
+
+
+@overload
+def stream_reader(
+    types: StreamTypes[ItemT, HeaderT, FooterT], stream: AsyncIterator[bytes]
+) -> StreamReaderFull[ItemT, HeaderT, FooterT]: ...
+
+
+def stream_reader(
+    types: StreamTypes[ItemT, HeaderTT, FooterTT], stream: AsyncIterator[bytes]
+) -> _StreamReader:
+    if types.header_type is None and types.footer_type is None:
+        return _StreamReader(types, stream)
+    if types.header_type is None:
+        return StreamReaderWithFooter(types, stream)
+    if types.footer_type is None:
+        return StreamReaderWithHeader(types, stream)
+
+    return StreamReaderFull(types, stream)
+
+
+# Writing ##############################################################################
+
+
+class _StreamWriterProtocol(Protocol[ItemT, HeaderTT, FooterTT]):
+    _stream_types: StreamTypes[ItemT, HeaderTT, FooterTT]
+    _last_sent: _Delimiter
+
+    def _serialize(self, obj: pydantic.BaseModel, delimiter: _Delimiter) -> bytes: ...
+
+
+class _StreamWriter(_Streamer[ItemT, HeaderTT, FooterTT]):
+    def __init__(self, types: StreamTypes[ItemT, HeaderTT, FooterTT]) -> None:
+        super().__init__(types)
+        self._last_sent = _Delimiter.NOT_SET
+        self._stream_types = types
+
+    @staticmethod
+    def _pack_tag(delimiter: _Delimiter, length: int) -> bytes:
+        return struct.pack(">BI", delimiter.value, length)
+
+    def _serialize(self, obj: pydantic.BaseModel, delimiter: _Delimiter) -> bytes:
+        data_bytes = obj.model_dump_json().encode()
+        data = bytearray(self._pack_tag(delimiter, len(data_bytes)))
+        data.extend(data_bytes)
+        # Starlette cannot handle byte array, but view works..
+        return memoryview(data)
+
+    def yield_item(self, item: ItemT) -> bytes:
+        if self._last_sent in (_Delimiter.FOOTER, _Delimiter.END):
+            raise ValueError("Cannot yield item after sending footer / closing stream.")
+        self._last_sent = _Delimiter.ITEM
+        return self._serialize(item, _Delimiter.ITEM)
+
+
+class _HeaderWriteMixin(_Streamer[ItemT, HeaderT, FooterTT]):
+    def yield_header(
+        self: _StreamWriterProtocol[ItemT, HeaderT, FooterTT], header: HeaderT
+    ) -> bytes:
+        if self._last_sent != _Delimiter.NOT_SET:
+            raise ValueError("Cannot yield header after other data has been sent.")
+        self._last_sent = _Delimiter.HEADER
+        return self._serialize(header, _Delimiter.HEADER)
+
+
+class _FooterWriteMixin(_Streamer[ItemT, HeaderTT, FooterT]):
+    def yield_footer(
+        self: _StreamWriterProtocol[ItemT, HeaderTT, FooterT], footer: FooterT
+    ) -> bytes:
+        if self._last_sent == _Delimiter.END:
+            raise ValueError("Cannot yield footer after closing stream.")
+        self._last_sent = _Delimiter.FOOTER
+        return self._serialize(footer, _Delimiter.FOOTER)
+
+
+class StreamWriterWithHeader(
+    _StreamWriter[ItemT, HeaderT, FooterTT], _HeaderWriteMixin[ItemT, HeaderT, FooterTT]
+): ...
+
+
+class StreamWriterWithFooter(
+    _StreamWriter[ItemT, HeaderTT, FooterT], _FooterWriteMixin[ItemT, HeaderTT, FooterT]
+): ...
+
+
+class StreamWriterFull(
+    _StreamWriter[ItemT, HeaderT, FooterT],
+    _HeaderWriteMixin[ItemT, HeaderT, FooterT],
+    _FooterWriteMixin[ItemT, HeaderT, FooterT],
+): ...
+
+
+@overload
+def stream_writer(
+    types: StreamTypes[ItemT, None, None],
+) -> _StreamWriter[ItemT, None, None]: ...
+
+
+@overload
+def stream_writer(
+    types: StreamTypes[ItemT, HeaderT, None],
+) -> StreamWriterWithHeader[ItemT, HeaderT, None]: ...
+
+
+@overload
+def stream_writer(
+    types: StreamTypes[ItemT, None, FooterT],
+) -> StreamWriterWithFooter[ItemT, None, FooterT]: ...
+
+
+@overload
+def stream_writer(
+    types: StreamTypes[ItemT, HeaderT, FooterT],
+) -> StreamWriterFull[ItemT, HeaderT, FooterT]: ...
+
+
+def stream_writer(types: StreamTypes[ItemT, HeaderTT, FooterTT]) -> _StreamWriter:
+    if types.header_type is None and types.footer_type is None:
+        return _StreamWriter(types)
+    if types.header_type is None:
+        return StreamWriterWithFooter(types)
+    if types.footer_type is None:
+        return StreamWriterWithHeader(types)
+
+    return StreamWriterFull(types)