Annotations for state_deltas.py (#11316)

I was sad that I couldn't do better for
`_curr_state_delta_stream_cache`. At least it's explicitly called out in
a comment with #TODO.

1 files changed, 13 insertions, 3 deletions

diff --git a/synapse/storage/databases/main/state_deltas.py b/synapse/storage/databases/main/state_deltas.py
index a89747d741..7f3624b128 100644
--- a/synapse/storage/databases/main/state_deltas.py
+++ b/synapse/storage/databases/main/state_deltas.py
@@ -16,11 +16,17 @@ import logging
 from typing import Any, Dict, List, Tuple
 
 from synapse.storage._base import SQLBaseStore
+from synapse.storage.database import LoggingTransaction
+from synapse.util.caches.stream_change_cache import StreamChangeCache
 
 logger = logging.getLogger(__name__)
 
 
 class StateDeltasStore(SQLBaseStore):
+    # This class must be mixed in with a child class which provides the following
+    # attribute. TODO: can we get static analysis to enforce this?
+    _curr_state_delta_stream_cache: StreamChangeCache
+
     async def get_current_state_deltas(
         self, prev_stream_id: int, max_stream_id: int
     ) -> Tuple[int, List[Dict[str, Any]]]:
@@ -60,7 +66,9 @@ class StateDeltasStore(SQLBaseStore):
             # max_stream_id.
             return max_stream_id, []
 
-        def get_current_state_deltas_txn(txn):
+        def get_current_state_deltas_txn(
+            txn: LoggingTransaction,
+        ) -> Tuple[int, List[Dict[str, Any]]]:
             # First we calculate the max stream id that will give us less than
             # N results.
             # We arbitrarily limit to 100 stream_id entries to ensure we don't
@@ -106,7 +114,9 @@ class StateDeltasStore(SQLBaseStore):
             "get_current_state_deltas", get_current_state_deltas_txn
         )
 
-    def _get_max_stream_id_in_current_state_deltas_txn(self, txn):
+    def _get_max_stream_id_in_current_state_deltas_txn(
+        self, txn: LoggingTransaction
+    ) -> int:
         return self.db_pool.simple_select_one_onecol_txn(
             txn,
             table="current_state_delta_stream",
@@ -114,7 +124,7 @@ class StateDeltasStore(SQLBaseStore):
             retcol="COALESCE(MAX(stream_id), -1)",
         )
 
-    async def get_max_stream_id_in_current_state_deltas(self):
+    async def get_max_stream_id_in_current_state_deltas(self) -> int:
         return await self.db_pool.runInteraction(
             "get_max_stream_id_in_current_state_deltas",
             self._get_max_stream_id_in_current_state_deltas_txn,