mypy: make quixstreams.state.* pass type checks

pyproject.toml

@@ -116,8 +116,6 @@ module = [
     "quixstreams.dataframe.*",
     "quixstreams.models.*",
     "quixstreams.platforms.*",
-    "quixstreams.state.*",
     "quixstreams.rowproducer.*"
 ]
 ignore_errors = true
-

quixstreams/dataframe/dataframe.py

@@ -8,7 +8,6 @@
 from typing import (
     Any,
     Callable,
-    Collection,
     Dict,
     List,
     Literal,
@@ -36,7 +35,7 @@
     VoidExecutor,
 )
 from quixstreams.models import (
-    HeaderValue,
+    HeadersTuples,
     MessageContext,
     Row,
     Topic,
@@ -626,8 +625,8 @@ def _set_timestamp_callback(
     def set_headers(
         self,
         func: Callable[
-            [Any, Any, int, List[Tuple[str, HeaderValue]]],
-            Collection[Tuple[str, HeaderValue]],
+            [Any, Any, int, HeadersTuples],
+            HeadersTuples,
         ],
     ) -> Self:
         """
@@ -663,8 +662,8 @@ def _set_headers_callback(
             value: Any,
             key: Any,
             timestamp: int,
-            headers: Collection[Tuple[str, HeaderValue]],
-        ) -> Tuple[Any, Any, int, Collection[Tuple[str, HeaderValue]]]:
+            headers: HeadersTuples,
+        ) -> Tuple[Any, Any, int, HeadersTuples]:
             # Create a shallow copy of original headers to prevent potential mutations
             # of the same collection
             headers = list(headers) if headers else []
@@ -1102,7 +1101,7 @@ def sink(self, sink: BaseSink):
         self._processing_context.sink_manager.register(sink)
 
         def _sink_callback(
-            value: Any, key: Any, timestamp: int, headers: List[Tuple[str, HeaderValue]]
+            value: Any, key: Any, timestamp: int, headers: HeadersTuples
         ):
             ctx = message_context()
             sink.add(

quixstreams/models/messages.py

@@ -12,7 +12,7 @@ class KafkaMessage:
 
     def __init__(
         self,
-        key: Optional[MessageKey],
+        key: MessageKey,
         value: Optional[MessageValue],
         headers: Optional[Headers],
         timestamp: Optional[int] = None,

quixstreams/models/rows.py

@@ -1,7 +1,7 @@
 from typing import Any, Optional
 
 from .messagecontext import MessageContext
-from .types import MessageHeadersTuples
+from .types import KafkaHeaders
 
 
 class Row:
@@ -19,7 +19,7 @@ def __init__(
         key: Optional[Any],
         timestamp: int,
         context: MessageContext,
-        headers: Optional[MessageHeadersTuples] = None,
+        headers: KafkaHeaders = None,
     ):
         self.value = value
         self.key = key

quixstreams/models/serializers/base.py

@@ -1,5 +1,5 @@
 import abc
-from typing import Any, Optional, Union
+from typing import Any, Union
 
 from confluent_kafka.serialization import (
     MessageField,
@@ -9,7 +9,7 @@
 )
 from typing_extensions import Literal, TypeAlias
 
-from ..types import MessageHeadersMapping, MessageHeadersTuples
+from ..types import HeadersMapping, KafkaHeaders
 
 __all__ = (
     "SerializationContext",
@@ -34,7 +34,7 @@ def __init__(
         self,
         topic: str,
         field: MessageField,
-        headers: Optional[MessageHeadersTuples] = None,
+        headers: KafkaHeaders = None,
     ) -> None:
         self.topic = topic
         self.field = field
@@ -65,7 +65,7 @@ class Serializer(abc.ABC):
     """
 
     @property
-    def extra_headers(self) -> MessageHeadersMapping:
+    def extra_headers(self) -> HeadersMapping:
         """
         Informs producer to set additional headers
         for the message it will be serializing

quixstreams/models/serializers/quix.py

@@ -2,7 +2,7 @@
 import gzip
 from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Tuple, Union
 
-from quixstreams.models.types import MessageHeadersMapping
+from quixstreams.models.types import HeadersMapping
 from quixstreams.utils.json import (
     dumps as default_dumps,
 )
@@ -290,7 +290,7 @@ def __init__(
         self._as_legacy = as_legacy
 
     @property
-    def extra_headers(self) -> MessageHeadersMapping:
+    def extra_headers(self) -> HeadersMapping:
         # Legacy-formatted messages should not contain any headers
         return self._extra_headers if not self._as_legacy else {}
 

quixstreams/models/topics/topic.py

@@ -25,15 +25,15 @@
 from quixstreams.models.types import (
     ConfluentKafkaMessageProto,
     Headers,
-    MessageHeadersTuples,
+    KafkaHeaders,
 )
 
 __all__ = ("Topic", "TopicConfig", "TimestampExtractor")
 
 logger = logging.getLogger(__name__)
 
 TimestampExtractor = Callable[
-    [Any, Optional[MessageHeadersTuples], int, TimestampType],
+    [Any, KafkaHeaders, int, TimestampType],
     int,
 ]
 

quixstreams/models/topics/utils.py

@@ -1,13 +1,13 @@
-from typing import Optional
-
-from quixstreams.models.types import MessageHeadersMapping, MessageHeadersTuples
+from quixstreams.models.types import (
+    HeadersMapping,
+    HeadersTuples,
+    KafkaHeaders,
+)
 
 __all__ = ("merge_headers",)
 
 
-def merge_headers(
-    original: Optional[MessageHeadersTuples], other: MessageHeadersMapping
-) -> MessageHeadersTuples:
+def merge_headers(original: KafkaHeaders, other: HeadersMapping) -> HeadersTuples:
     """
     Merge two sets of Kafka message headers, overwriting headers in "origin"
     by the values from "other".

quixstreams/models/types.py

@@ -1,13 +1,16 @@
-from typing import Dict, List, Optional, Tuple, Union
+from typing import List, Mapping, Optional, Sequence, Tuple, Union
 
 from typing_extensions import Protocol
 
-MessageKey = Union[str, bytes]
+MessageKey = Optional[Union[str, bytes]]
 MessageValue = Union[str, bytes]
-HeaderValue = Optional[Union[str, bytes]]
-MessageHeadersTuples = List[Tuple[str, HeaderValue]]
-MessageHeadersMapping = Dict[str, HeaderValue]
-Headers = Union[MessageHeadersTuples, MessageHeadersMapping]
+
+HeadersValue = Union[str, bytes]
+HeadersMapping = Mapping[str, HeadersValue]
+HeadersTuples = Sequence[Tuple[str, HeadersValue]]
+Headers = Union[HeadersTuples, HeadersMapping]
+
+KafkaHeaders = Optional[List[Tuple[str, bytes]]]
 
 
 class ConfluentKafkaMessageProto(Protocol):
@@ -21,13 +24,13 @@ class ConfluentKafkaMessageProto(Protocol):
 
     """
 
-    def headers(self, *args, **kwargs) -> Optional[MessageHeadersTuples]: ...
+    def headers(self, *args, **kwargs) -> KafkaHeaders: ...
 
-    def key(self, *args, **kwargs) -> Optional[MessageKey]: ...
+    def key(self, *args, **kwargs) -> MessageKey: ...
 
     def offset(self, *args, **kwargs) -> int: ...
 
-    def partition(self, *args, **kwargs) -> int: ...
+    def partition(self, *args, **kwargs) -> Optional[int]: ...
 
     def timestamp(self, *args, **kwargs) -> Tuple[int, int]: ...
 

quixstreams/sinks/base/batch.py

@@ -1,8 +1,8 @@
 from collections import deque
 from itertools import islice
-from typing import Any, Deque, Iterable, Iterator, List, Tuple
+from typing import Any, Deque, Iterable, Iterator
 
-from quixstreams.models import HeaderValue
+from quixstreams.models import HeadersTuples
 
 from .item import SinkItem
 
@@ -47,7 +47,7 @@ def append(
         value: Any,
         key: Any,
         timestamp: int,
-        headers: List[Tuple[str, HeaderValue]],
+        headers: HeadersTuples,
         offset: int,
     ):
         self._buffer.append(

quixstreams/sinks/base/item.py

@@ -1,6 +1,6 @@
-from typing import Any, List, Tuple
+from typing import Any
 
-from quixstreams.models import HeaderValue
+from quixstreams.models import HeadersTuples
 
 __all__ = ("SinkItem",)
 
@@ -19,7 +19,7 @@ def __init__(
         value: Any,
         key: Any,
         timestamp: int,
-        headers: List[Tuple[str, HeaderValue]],
+        headers: HeadersTuples,
         offset: int,
     ):
         self.key = key

quixstreams/sinks/base/sink.py

@@ -1,8 +1,8 @@
 import abc
 import logging
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, Tuple
 
-from quixstreams.models import HeaderValue
+from quixstreams.models import HeadersTuples
 from quixstreams.sinks.base.batch import SinkBatch
 
 logger = logging.getLogger(__name__)
@@ -35,7 +35,7 @@ def add(
         value: Any,
         key: Any,
         timestamp: int,
-        headers: List[Tuple[str, HeaderValue]],
+        headers: HeadersTuples,
         topic: str,
         partition: int,
         offset: int,
@@ -95,7 +95,7 @@ def add(
         value: Any,
         key: Any,
         timestamp: int,
-        headers: List[Tuple[str, HeaderValue]],
+        headers: HeadersTuples,
         topic: str,
         partition: int,
         offset: int,

quixstreams/sinks/community/bigquery.py

@@ -18,7 +18,7 @@
     ) from exc
 
 from quixstreams.exceptions import QuixException
-from quixstreams.models import HeaderValue
+from quixstreams.models import HeadersTuples
 from quixstreams.sinks import BatchingSink, SinkBatch
 
 __all__ = ("BigQuerySink", "BigQuerySinkException")
@@ -168,7 +168,7 @@ def add(
         value: Any,
         key: Any,
         timestamp: int,
-        headers: list[tuple[str, HeaderValue]],
+        headers: HeadersTuples,
         topic: str,
         partition: int,
         offset: int,

quixstreams/sinks/community/kinesis.py

@@ -13,7 +13,7 @@
         'run "pip install quixstreams[kinesis]" to use KinesisSink'
     ) from exc
 
-from quixstreams.models.types import HeaderValue
+from quixstreams.models.types import HeadersTuples
 from quixstreams.sinks.base import BaseSink
 from quixstreams.sinks.base.exceptions import SinkBackpressureError
 
@@ -82,7 +82,7 @@ def add(
         value: Any,
         key: Any,
         timestamp: int,
-        headers: list[tuple[str, HeaderValue]],
+        headers: HeadersTuples,
         topic: str,
         partition: int,
         offset: int,

quixstreams/sinks/community/postgresql.py

@@ -14,7 +14,7 @@
     ) from exc
 
 from quixstreams.exceptions import QuixException
-from quixstreams.models import HeaderValue
+from quixstreams.models import HeadersTuples
 from quixstreams.sinks import BatchingSink, SinkBatch
 
 __all__ = ("PostgreSQLSink", "PostgreSQLSinkException")
@@ -115,7 +115,7 @@ def add(
         value: Any,
         key: Any,
         timestamp: int,
-        headers: list[tuple[str, HeaderValue]],
+        headers: HeadersTuples,
         topic: str,
         partition: int,
         offset: int,

quixstreams/sinks/community/pubsub.py

@@ -14,7 +14,7 @@
         'run "pip install quixstreams[pubsub]" to use PubSubSink'
     ) from exc
 
-from quixstreams.models.types import HeaderValue
+from quixstreams.models.types import HeadersTuples
 from quixstreams.sinks.base import BaseSink, SinkBackpressureError
 
 __all__ = ("PubSubSink", "PubSubTopicNotFoundError")
@@ -83,7 +83,7 @@ def add(
         value: Any,
         key: Any,
         timestamp: int,
-        headers: list[tuple[str, HeaderValue]],
+        headers: HeadersTuples,
         topic: str,
         partition: int,
         offset: int,

quixstreams/sinks/core/influxdb3.py

@@ -1,9 +1,9 @@
 import logging
 import sys
 import time
-from typing import Any, Iterable, List, Mapping, Optional, Tuple
+from typing import Any, Iterable, Mapping, Optional
 
-from quixstreams.models import HeaderValue
+from quixstreams.models import HeadersTuples
 
 try:
     import influxdb_client_3
@@ -127,7 +127,7 @@ def add(
         value: Any,
         key: Any,
         timestamp: int,
-        headers: List[Tuple[str, HeaderValue]],
+        headers: HeadersTuples,
         topic: str,
         partition: int,
         offset: int,