Fix type issues

2 files changed, 11 insertions, 8 deletions

diff --git a/synapse/storage/databases/state/store.py b/synapse/storage/databases/state/store.py
index 04eddbeffd..dbe152fc58 100644
--- a/synapse/storage/databases/state/store.py
+++ b/synapse/storage/databases/state/store.py
@@ -38,6 +38,12 @@ logger = logging.getLogger(__name__)
 # XXX
 UNKNOWN = Any  # TODO
 
+
+InflightStateGroupCacheKey = Union[
+    Tuple[int, StateFilter], Tuple[int, str, Optional[str]]
+]
+
+
 MAX_STATE_DELTA_HOPS = 100
 
 
@@ -99,9 +105,8 @@ class StateGroupDataStore(StateBackgroundUpdateStore, SQLBaseStore):
             500000,
         )
 
-        # XXX ADD TYPE
         self._state_group_inflight_cache: MultiKeyResponseCache[
-            ...
+            InflightStateGroupCacheKey, Dict[int, StateMap[str]]
         ] = MultiKeyResponseCache(
             self.hs.get_clock(),
             "*stateGroupInflightCache*",
@@ -418,9 +423,7 @@ class StateGroupDataStore(StateBackgroundUpdateStore, SQLBaseStore):
             # make a list of keys for us to store in the in-flight cache
             # this should list all the keys that the request will pick up from
             # the database.
-            keys: List[
-                Union[Tuple[int, StateFilter], Tuple[int, str, Optional[str]]]
-            ] = []
+            keys: List[InflightStateGroupCacheKey] = []
             for group in inflight_cache_misses:
                 if db_state_filter.include_others:
                     # We can't properly add cache keys for all the 'other'
diff --git a/synapse/util/caches/multi_key_response_cache.py b/synapse/util/caches/multi_key_response_cache.py
index c87bb3734e..3bf7d3f4c7 100644
--- a/synapse/util/caches/multi_key_response_cache.py
+++ b/synapse/util/caches/multi_key_response_cache.py
@@ -54,7 +54,7 @@ class MultiKeyResponseCacheContext(Generic[KV]):
     """
 
 
-class MultiKeyResponseCache(Generic[KV]):
+class MultiKeyResponseCache(Generic[KV, RV]):
     """
     This caches a deferred response. Until the deferred completes it will be
     returned from the cache. This means that if the client retries the request
@@ -69,7 +69,7 @@ class MultiKeyResponseCache(Generic[KV]):
         # This is poorly-named: it includes both complete and incomplete results.
         # We keep complete results rather than switching to absolute values because
         # that makes it easier to cache Failure results.
-        self.pending_result_cache: Dict[KV, ObservableDeferred] = {}
+        self.pending_result_cache: Dict[KV, ObservableDeferred[RV]] = {}
 
         self.clock = clock
         self.timeout_sec = timeout_ms / 1000.0
@@ -109,7 +109,7 @@ class MultiKeyResponseCache(Generic[KV]):
             return None
 
     def _set(
-        self, context: MultiKeyResponseCacheContext[KV], deferred: defer.Deferred
+        self, context: MultiKeyResponseCacheContext[KV], deferred: "defer.Deferred[RV]"
     ) -> "defer.Deferred[RV]":
         """Set the entry for the given key to the given deferred.