typing: make state protocols and ABCs generic (#777)

* typing: make state protocols and ABCs generic

* windows: set transaction type in count windows

Author: ovv

quixstreams/dataframe/windows/count_based.py

@@ -58,7 +58,7 @@ def process_window(
         value: Any,
         key: Any,
         timestamp_ms: int,
-        transaction: WindowedPartitionTransaction,
+        transaction: WindowedPartitionTransaction[str, CountWindowsData],
     ) -> tuple[Iterable[WindowKeyResult], Iterable[WindowKeyResult]]:
         """
         Count based windows are different from time based windows as we don't
@@ -81,9 +81,7 @@ def process_window(
         optimisation. Instead the msg id reset to 0 on every new window.
         """
         state = transaction.as_state(prefix=key)
-        data = state.get(key=self.STATE_KEY)
-        if data is None:
-            data = CountWindowsData(windows=[])
+        data = state.get(key=self.STATE_KEY, default=CountWindowsData(windows=[]))
 
         msg_id = None
         if len(data["windows"]) == 0:

quixstreams/state/base/state.py

@@ -1,6 +1,6 @@
 import logging
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Generic, Optional, TypeVar, overload
 
 if TYPE_CHECKING:
     from .transaction import PartitionTransaction
@@ -10,13 +10,23 @@
 logger = logging.getLogger(__name__)
 
 
-class State(ABC):
+K = TypeVar("K")
+V = TypeVar("V")
+
+
+class State(ABC, Generic[K, V]):
     """
     Primary interface for working with key-value state data from `StreamingDataFrame`
     """
 
+    @overload
+    def get(self, key: K) -> Optional[V]: ...
+
+    @overload
+    def get(self, key: K, default: V) -> V: ...
+
     @abstractmethod
-    def get(self, key: Any, default: Any = None) -> Optional[Any]:
+    def get(self, key: K, default: Optional[V] = None) -> Optional[V]:
         """
         Get the value for key if key is present in the state, else default
 
@@ -27,7 +37,7 @@ def get(self, key: Any, default: Any = None) -> Optional[Any]:
         ...
 
     @abstractmethod
-    def set(self, key: Any, value: Any):
+    def set(self, key: K, value: V):
         """
         Set value for the key.
         :param key: key
@@ -36,7 +46,7 @@ def set(self, key: Any, value: Any):
         ...
 
     @abstractmethod
-    def delete(self, key: Any):
+    def delete(self, key: K):
         """
         Delete value for the key.
 
@@ -46,7 +56,7 @@ def delete(self, key: Any):
         ...
 
     @abstractmethod
-    def exists(self, key: Any) -> bool:
+    def exists(self, key: K) -> bool:
         """
         Check if the key exists in state.
         :param key: key
@@ -70,7 +80,13 @@ def __init__(self, prefix: bytes, transaction: "PartitionTransaction"):
         self._prefix = prefix
         self._transaction = transaction
 
-    def get(self, key: Any, default: Any = None) -> Optional[Any]:
+    @overload
+    def get(self, key: K) -> Optional[V]: ...
+
+    @overload
+    def get(self, key: K, default: V) -> V: ...
+
+    def get(self, key: K, default: Optional[V] = None) -> Optional[V]:
         """
         Get the value for key if key is present in the state, else default
 
@@ -80,15 +96,15 @@ def get(self, key: Any, default: Any = None) -> Optional[Any]:
         """
         return self._transaction.get(key=key, prefix=self._prefix, default=default)
 
-    def set(self, key: Any, value: Any):
+    def set(self, key: K, value: V):
         """
         Set value for the key.
         :param key: key
         :param value: value
         """
         return self._transaction.set(key=key, value=value, prefix=self._prefix)
 
-    def delete(self, key: Any):
+    def delete(self, key: K):
         """
         Delete value for the key.
 
@@ -97,7 +113,7 @@ def delete(self, key: Any):
         """
         return self._transaction.delete(key=key, prefix=self._prefix)
 
-    def exists(self, key: Any) -> bool:
+    def exists(self, key: K) -> bool:
         """
         Check if the key exists in state.
         :param key: key

quixstreams/state/base/transaction.py

@@ -7,10 +7,13 @@
     TYPE_CHECKING,
     Any,
     Dict,
+    Generic,
     Optional,
     Set,
     Tuple,
+    TypeVar,
     Union,
+    overload,
 )
 
 from quixstreams.models import Headers
@@ -179,7 +182,11 @@ def _wrapper(tx: "PartitionTransaction", *args, **kwargs):
     return wrapper
 
 
-class PartitionTransaction(ABC):
+K = TypeVar("K")
+V = TypeVar("V")
+
+
+class PartitionTransaction(ABC, Generic[K, V]):
     """
     A transaction class to perform simple key-value operations like
     "get", "set", "delete" and "exists" on a single storage partition.
@@ -257,18 +264,18 @@ def changelog_topic_partition(self) -> Optional[Tuple[str, int]]:
             self.changelog_producer.partition,
         )
 
-    def _serialize_value(self, value: Any) -> bytes:
+    def _serialize_value(self, value: V) -> bytes:
         return serialize(value, dumps=self._dumps)
 
-    def _deserialize_value(self, value: bytes) -> Any:
+    def _deserialize_value(self, value: bytes) -> V:
         return deserialize(value, loads=self._loads)
 
-    def _serialize_key(self, key: Any, prefix: bytes) -> bytes:
+    def _serialize_key(self, key: K, prefix: bytes) -> bytes:
         key_bytes = serialize(key, dumps=self._dumps)
         prefix = prefix + SEPARATOR if prefix else b""
         return prefix + key_bytes
 
-    def as_state(self, prefix: Any = DEFAULT_PREFIX) -> State:
+    def as_state(self, prefix: Any = DEFAULT_PREFIX) -> State[K, V]:
         """
         Create an instance implementing the `State` protocol to be provided
         to `StreamingDataFrame` functions.
@@ -286,14 +293,20 @@ def as_state(self, prefix: Any = DEFAULT_PREFIX) -> State:
             ),
         )
 
+    @overload
+    def get(self, key: K, prefix: bytes, cf_name: str = "default") -> Optional[V]: ...
+
+    @overload
+    def get(self, key: K, prefix: bytes, default: V, cf_name: str = "default") -> V: ...
+
     @validate_transaction_status(PartitionTransactionStatus.STARTED)
     def get(
         self,
-        key: Any,
+        key: K,
         prefix: bytes,
-        default: Any = None,
+        default: Optional[V] = None,
         cf_name: str = "default",
-    ) -> Any:
+    ) -> Optional[V]:
         """
         Get a key from the store.
 
@@ -323,7 +336,7 @@ def get(
         return self._deserialize_value(stored)
 
     @validate_transaction_status(PartitionTransactionStatus.STARTED)
-    def set(self, key: Any, value: Any, prefix: bytes, cf_name: str = "default"):
+    def set(self, key: K, value: V, prefix: bytes, cf_name: str = "default"):
         """
         Set value for the key.
         :param key: key
@@ -346,7 +359,7 @@ def set(self, key: Any, value: Any, prefix: bytes, cf_name: str = "default"):
             raise
 
     @validate_transaction_status(PartitionTransactionStatus.STARTED)
-    def delete(self, key: Any, prefix: bytes, cf_name: str = "default"):
+    def delete(self, key: K, prefix: bytes, cf_name: str = "default"):
         """
         Delete value for the key.
 
@@ -365,7 +378,7 @@ def delete(self, key: Any, prefix: bytes, cf_name: str = "default"):
             raise
 
     @validate_transaction_status(PartitionTransactionStatus.STARTED)
-    def exists(self, key: Any, prefix: bytes, cf_name: str = "default") -> bool:
+    def exists(self, key: K, prefix: bytes, cf_name: str = "default") -> bool:
         """
         Check if the key exists in state.
         :param key: key

quixstreams/state/rocksdb/windowed/transaction.py

@@ -358,13 +358,14 @@ def expire_all_windows(
                     prefix, start, end = parse_window_key(key)
                     to_delete.add((prefix, start, end))
                     if collect:
-                        value = self.get_from_collection(
+                        value: Any = self.get_from_collection(
                             start=start,
                             end=end,
                             prefix=prefix,
                         )
                     else:
                         value = self.get(encode_integer_pair(start, end), prefix=prefix)
+                        assert value is not None  # noqa: S101
 
                     yield (start, end), value, prefix
         else:
@@ -385,6 +386,7 @@ def expire_all_windows(
                         )
                     else:
                         value = self.get(encode_integer_pair(start, end), prefix=prefix)
+                        assert value is not None  # noqa: S101
 
                     yield (start, end), value, prefix
 
@@ -603,14 +605,15 @@ def _get_next_count(self) -> int:
         :return: Next sequential counter value
         """
         cache = self._global_counter
-        kwargs = {"key": cache.key, "prefix": b"", "cf_name": cache.cf_name}
 
         if cache.counter is None:
-            cache.counter = self.get(default=-1, **kwargs)
+            cache.counter = self.get(
+                default=-1, key=cache.key, prefix=b"", cf_name=cache.cf_name
+            )
 
         cache.counter += 1
 
-        self.set(value=cache.counter, **kwargs)
+        self.set(value=cache.counter, key=cache.key, prefix=b"", cf_name=cache.cf_name)
         return cache.counter