typing (#876)

Author: ovv

quixstreams/state/base/partition.py

@@ -1,6 +1,6 @@
 import logging
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Generic, Literal, Optional, TypeVar, Union, cast
+from typing import TYPE_CHECKING, Literal, Optional, Type, Union
 
 from quixstreams.state.metadata import (
     Marker,
@@ -17,17 +17,15 @@
 
 logger = logging.getLogger(__name__)
 
-T = TypeVar("T", bound=PartitionTransaction)
 
-
-class StorePartition(ABC, Generic[T]):
+class StorePartition(ABC):
     """
     A base class to access state in the underlying storage.
     It represents a single instance of some storage (e.g. a single database for
     the persistent storage).
     """
 
-    partition_transaction_class = PartitionTransaction
+    partition_transaction_class: Type[PartitionTransaction] = PartitionTransaction
 
     def __init__(
         self,
@@ -110,20 +108,17 @@ def recover_from_changelog_message(
         :param offset: changelog message offset
         """
 
-    def begin(self) -> T:
+    def begin(self) -> PartitionTransaction:
         """
         Start a new `PartitionTransaction`
 
         Using `PartitionTransaction` is a recommended way for accessing the data.
         """
-        return cast(
-            T,
-            self.partition_transaction_class(
-                partition=self,
-                dumps=self._dumps,
-                loads=self._loads,
-                changelog_producer=self._changelog_producer,
-            ),
+        return self.partition_transaction_class(
+            partition=self,
+            dumps=self._dumps,
+            loads=self._loads,
+            changelog_producer=self._changelog_producer,
         )
 
     def __enter__(self):

quixstreams/state/rocksdb/partition.py

@@ -6,15 +6,13 @@
     List,
     Literal,
     Optional,
-    TypeVar,
     Union,
     cast,
 )
 
 from rocksdict import AccessType, ColumnFamily, Rdict, ReadOptions, WriteBatch
 
 from quixstreams.state.base import (
-    PartitionTransaction,
     PartitionTransactionCache,
     StorePartition,
 )
@@ -35,10 +33,8 @@
 
 logger = logging.getLogger(__name__)
 
-T = TypeVar("T", bound=PartitionTransaction)
 
-
-class RocksDBStorePartition(StorePartition[T]):
+class RocksDBStorePartition(StorePartition):
     """
     A base class to access state in RocksDB.
     It represents a single RocksDB database.

quixstreams/state/rocksdb/windowed/partition.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Iterator, cast
+from typing import Iterator, Type, cast
 
 from ..partition import RocksDBStorePartition
 from .metadata import (
@@ -15,9 +15,7 @@
 logger = logging.getLogger(__name__)
 
 
-class WindowedRocksDBStorePartition(
-    RocksDBStorePartition[WindowedRocksDBPartitionTransaction]
-):
+class WindowedRocksDBStorePartition(RocksDBStorePartition):
     """
     A base class to access windowed state in RocksDB.
     It represents a single RocksDB database.
@@ -26,7 +24,9 @@ class WindowedRocksDBStorePartition(
     stores the expiration index to delete expired windows.
     """
 
-    partition_transaction_class = WindowedRocksDBPartitionTransaction
+    partition_transaction_class: Type[WindowedRocksDBPartitionTransaction] = (
+        WindowedRocksDBPartitionTransaction
+    )
     additional_column_families = (
         LATEST_DELETED_VALUE_CF_NAME,
         LATEST_EXPIRED_WINDOW_CF_NAME,
@@ -47,3 +47,16 @@ def iter_keys(self, cf_name: str = "default") -> Iterator[bytes]:
         """
         cf_dict = self.get_column_family(cf_name)
         return cast(Iterator[bytes], cf_dict.keys())
+
+    def begin(self) -> WindowedRocksDBPartitionTransaction:
+        """
+        Start a new `WindowedRocksDBPartitionTransaction`
+
+        Using `WindowedRocksDBPartitionTransaction` is a recommended way for accessing the data.
+        """
+        return self.partition_transaction_class(
+            partition=self,
+            dumps=self._dumps,
+            loads=self._loads,
+            changelog_producer=self._changelog_producer,
+        )

quixstreams/state/rocksdb/windowed/transaction.py

@@ -41,7 +41,7 @@
     from .partition import WindowedRocksDBStorePartition
 
 
-class WindowedRocksDBPartitionTransaction(PartitionTransaction):
+class WindowedRocksDBPartitionTransaction(PartitionTransaction[bytes, dict]):
     def __init__(
         self,
         partition: "WindowedRocksDBStorePartition",
@@ -469,8 +469,7 @@ def get_windows(
         ):
             _, start, end = parse_window_key(key)
             if start_from_ms < start <= start_to_ms:
-                value = self._deserialize_value(value)
-                result.append(((start, end), value, prefix))
+                result.append(((start, end), self._deserialize_value(value), prefix))
 
         return result
 
@@ -532,7 +531,7 @@ def _get_items(
 
     def _get_timestamp(
         self, cache: TimestampsCache, prefix: bytes, default: Any = None
-    ) -> Any:
+    ) -> int:
         cached_ts = cache.timestamps.get(prefix)
         if cached_ts is not None:
             return cached_ts
@@ -585,7 +584,7 @@ def _get_next_count(self) -> int:
         if cache.counter is None:
             cache.counter = self.get(
                 default=-1, key=cache.key, prefix=b"", cf_name=cache.cf_name
-            )
+            )  # type:ignore[call-overload]
 
         cache.counter += 1