Add generic typing to the StorePartition classes to correctly reveal the transaction type

Author: daniil-quix

quixstreams/state/base/partition.py

@@ -1,6 +1,6 @@
 import logging
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Literal, Optional, Union
+from typing import TYPE_CHECKING, Generic, Literal, Optional, TypeVar, Union, cast
 
 from quixstreams.state.metadata import (
     Marker,
@@ -17,8 +17,10 @@
 
 logger = logging.getLogger(__name__)
 
+T = TypeVar("T", bound=PartitionTransaction)
 
-class StorePartition(ABC):
+
+class StorePartition(ABC, Generic[T]):
     """
     A base class to access state in the underlying storage.
     It represents a single instance of some storage (e.g. a single database for
@@ -108,17 +110,20 @@ def recover_from_changelog_message(
         :param offset: changelog message offset
         """
 
-    def begin(self) -> PartitionTransaction:
+    def begin(self) -> T:
         """
         Start a new `PartitionTransaction`
 
         Using `PartitionTransaction` 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,
+        return cast(
+            T,
+            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

@@ -1,17 +1,27 @@
 import logging
 import time
-from typing import Dict, Iterator, List, Literal, Optional, Union, cast
+from typing import (
+    Dict,
+    Iterator,
+    List,
+    Literal,
+    Optional,
+    TypeVar,
+    Union,
+    cast,
+)
 
 from rocksdict import AccessType, ColumnFamily, Rdict, ReadOptions, WriteBatch
 
-from quixstreams.state.base import PartitionTransactionCache, StorePartition
+from quixstreams.state.base import (
+    PartitionTransaction,
+    PartitionTransactionCache,
+    StorePartition,
+)
 from quixstreams.state.exceptions import ColumnFamilyDoesNotExist
 from quixstreams.state.metadata import METADATA_CF_NAME, Marker
 from quixstreams.state.recovery import ChangelogProducer
-from quixstreams.state.serialization import (
-    int_from_int64_bytes,
-    int_to_int64_bytes,
-)
+from quixstreams.state.serialization import int_from_int64_bytes, int_to_int64_bytes
 
 from .exceptions import ColumnFamilyAlreadyExists
 from .metadata import (
@@ -22,10 +32,13 @@
 
 __all__ = ("RocksDBStorePartition",)
 
+
 logger = logging.getLogger(__name__)
 
+T = TypeVar("T", bound=PartitionTransaction)
+
 
-class RocksDBStorePartition(StorePartition):
+class RocksDBStorePartition(StorePartition[T]):
     """
     A base class to access state in RocksDB.
     It represents a single RocksDB database.

quixstreams/state/rocksdb/windowed/partition.py

@@ -15,7 +15,9 @@
 logger = logging.getLogger(__name__)
 
 
-class WindowedRocksDBStorePartition(RocksDBStorePartition):
+class WindowedRocksDBStorePartition(
+    RocksDBStorePartition[WindowedRocksDBPartitionTransaction]
+):
     """
     A base class to access windowed state in RocksDB.
     It represents a single RocksDB database.