feature: Refine exposed protocol types

By refining the exposed protocol types to include unions types for the
different kinds of headers, we allow typing at call-sites to work
smoother. This also enables using exhaustiveness checking in a place
where it wasn't possible before, when matching the type of an entity's
header schema. See the replaced NotImplementedError in the diff for
reference.

codegen/generate_schema.py

@@ -54,6 +54,7 @@
 from dataclasses import dataclass, field
 from typing import Annotated, ClassVar
 import uuid
+from kio.static.base import ApiMessage
 from kio.static.primitive import i8
 from kio.static.primitive import i16
 from kio.static.primitive import i32
@@ -66,7 +67,6 @@
 from kio.static.primitive import i32Timedelta
 from kio.static.primitive import i64Timedelta
 from kio.static.primitive import TZAware
-from kio.static.protocol import ApiMessage
 from kio.static.constants import ErrorCode
 '''
 

codegen/generate_tests.py

@@ -9,7 +9,7 @@
 from types import ModuleType
 
 import kio.schema
-from kio.static.protocol import ApiMessage
+from kio.static.base import ApiMessage
 
 from .case import to_snake_case
 

src/kio/static/base.py

@@ -0,0 +1,5 @@
+__all__ = ("ApiMessage",)
+
+
+class ApiMessage:
+    pass

src/kio/static/protocol.py

@@ -1,15 +1,44 @@
+from __future__ import annotations
+
 from typing import ClassVar
 from typing import Protocol
+from typing import TypeAlias
 
+import kio.schema.request_header.v0
+import kio.schema.request_header.v1
+import kio.schema.request_header.v2
+import kio.schema.response_header.v0
+import kio.schema.response_header.v1
 from kio._utils import DataclassInstance
 
 from .primitive import i16
 
-__all__ = ("ApiMessage", "Entity", "Payload")
+__all__ = (
+    "Entity",
+    "Payload",
+    "ResponsePayload",
+    "RequestPayload",
+    "RequestHeader",
+    "ResponseHeader",
+    "Header",
+    "HeaderV0RequestPayload",
+    "HeaderV1RequestPayload",
+    "HeaderV2RequestPayload",
+    "HeaderV0ResponsePayload",
+    "HeaderV1ResponsePayload",
+)
 
 
-class ApiMessage:
-    pass
+RequestHeader: TypeAlias = (
+    kio.schema.request_header.v0.RequestHeader
+    | kio.schema.request_header.v1.RequestHeader
+    | kio.schema.request_header.v2.RequestHeader
+)
+ResponseHeader: TypeAlias = (
+    kio.schema.response_header.v0.ResponseHeader
+    | kio.schema.response_header.v1.ResponseHeader
+)
+Header: TypeAlias = RequestHeader | ResponseHeader
 
 
 class Entity(DataclassInstance, Protocol):
@@ -52,8 +81,36 @@ class Payload(DataclassInstance, Protocol):
     # required to accept ANY subtype of Entity as set-type. This is quirky but sound
     # from static type checking perspective.
     @property
-    def __header_schema__(self) -> type[Entity]:
+    def __header_schema__(self) -> type[Header]:
         """
         The header entity type that should be used when sending or receiving this
         payload.
         """
+
+
+class HeaderV0RequestPayload(Payload, Protocol):
+    __header_schema__: ClassVar[type[kio.schema.request_header.v0.RequestHeader]]
+
+
+class HeaderV1RequestPayload(Payload, Protocol):
+    __header_schema__: ClassVar[type[kio.schema.request_header.v1.RequestHeader]]
+
+
+class HeaderV2RequestPayload(Payload, Protocol):
+    __header_schema__: ClassVar[type[kio.schema.request_header.v2.RequestHeader]]
+
+
+RequestPayload: TypeAlias = (
+    HeaderV0RequestPayload | HeaderV1RequestPayload | HeaderV2RequestPayload
+)
+
+
+class HeaderV0ResponsePayload(Payload, Protocol):
+    __header_schema__: ClassVar[type[kio.schema.response_header.v0.ResponseHeader]]
+
+
+class HeaderV1ResponsePayload(Payload, Protocol):
+    __header_schema__: ClassVar[type[kio.schema.response_header.v1.ResponseHeader]]
+
+
+ResponsePayload: TypeAlias = HeaderV0ResponsePayload | HeaderV1ResponsePayload

tests/test_integration.py

@@ -8,6 +8,7 @@
 from typing import Any
 from typing import Final
 from typing import TypeVar
+from typing import assert_never
 from unittest import mock
 
 import pytest
@@ -45,8 +46,9 @@
 from kio.static.primitive import i16
 from kio.static.primitive import i32
 from kio.static.primitive import i32Timedelta
-from kio.static.protocol import Entity
-from kio.static.protocol import Payload
+from kio.static.protocol import RequestHeader
+from kio.static.protocol import RequestPayload
+from kio.static.protocol import ResponsePayload
 
 pytestmark = pytest.mark.integration
 
@@ -55,12 +57,12 @@
 
 def write_request_header(
     buffer: Writable,
-    payload: Payload,
+    payload: RequestPayload,
     correlation_id: i32,
     client_id: str | None,
 ) -> None:
     header_schema = payload.__header_schema__
-    header: Entity
+    header: RequestHeader
 
     if issubclass(header_schema, kio.schema.request_header.v0.header.RequestHeader):
         header = header_schema(
@@ -82,14 +84,14 @@ def write_request_header(
             client_id=client_id,
         )
     else:
-        raise NotImplementedError(f"Unknown request header schema: {header_schema}")
+        assert_never(header_schema)
 
     entity_writer(header_schema)(buffer, header)  # type: ignore[arg-type]
 
 
 async def send(
     stream: StreamWriter,
-    payload: Payload,
+    payload: RequestPayload,
     correlation_id: i32,
 ) -> None:
     write_request = entity_writer(type(payload))
@@ -126,7 +128,7 @@ async def read_response_bytes(stream: StreamReader) -> io.BytesIO:
     return io.BytesIO(await stream.read(response_length))
 
 
-R = TypeVar("R", bound=Payload)
+R = TypeVar("R", bound=ResponsePayload)
 
 
 def parse_response(
@@ -146,7 +148,7 @@ def parse_response(
 
 
 async def make_request(
-    request: Payload,
+    request: RequestPayload,
     response_type: type[R],
 ) -> R:
     correlation_id = i32(secrets.randbelow(i32.__high__ + 1))