Fix cache call signature to accept `on_invalidate`. (#8684)

Cached functions accept an `on_invalidate` function, which we failed to add to the type signature. It's rarely used in the files that we have typed, which is why we haven't noticed it before.

3 files changed, 35 insertions, 16 deletions

diff --git a/changelog.d/8684.misc b/changelog.d/8684.misc
new file mode 100644
index 0000000000..1d23d42926
--- /dev/null
+++ b/changelog.d/8684.misc
@@ -0,0 +1 @@
+Fix typing info on cache call signature to accept `on_invalidate`.
diff --git a/scripts-dev/mypy_synapse_plugin.py b/scripts-dev/mypy_synapse_plugin.py
index a5b88731f1..5882f3a0b0 100644
--- a/scripts-dev/mypy_synapse_plugin.py
+++ b/scripts-dev/mypy_synapse_plugin.py
@@ -19,9 +19,10 @@ can crop up, e.g the cache descriptors.
 
 from typing import Callable, Optional
 
+from mypy.nodes import ARG_NAMED_OPT
 from mypy.plugin import MethodSigContext, Plugin
 from mypy.typeops import bind_self
-from mypy.types import CallableType
+from mypy.types import CallableType, NoneType
 
 
 class SynapsePlugin(Plugin):
@@ -40,8 +41,9 @@ def cached_function_method_signature(ctx: MethodSigContext) -> CallableType:
 
     It already has *almost* the correct signature, except:
 
-        1. the `self` argument needs to be marked as "bound"; and
-        2. any `cache_context` argument should be removed.
+        1. the `self` argument needs to be marked as "bound";
+        2. any `cache_context` argument should be removed;
+        3. an optional keyword argument `on_invalidated` should be added.
     """
 
     # First we mark this as a bound function signature.
@@ -58,19 +60,33 @@ def cached_function_method_signature(ctx: MethodSigContext) -> CallableType:
             context_arg_index = idx
             break
 
+    arg_types = list(signature.arg_types)
+    arg_names = list(signature.arg_names)
+    arg_kinds = list(signature.arg_kinds)
+
     if context_arg_index:
-        arg_types = list(signature.arg_types)
         arg_types.pop(context_arg_index)
-
-        arg_names = list(signature.arg_names)
         arg_names.pop(context_arg_index)
-
-        arg_kinds = list(signature.arg_kinds)
         arg_kinds.pop(context_arg_index)
 
-        signature = signature.copy_modified(
-            arg_types=arg_types, arg_names=arg_names, arg_kinds=arg_kinds,
-        )
+    # Third, we add an optional "on_invalidate" argument.
+    #
+    # This is a callable which accepts no input and returns nothing.
+    calltyp = CallableType(
+        arg_types=[],
+        arg_kinds=[],
+        arg_names=[],
+        ret_type=NoneType(),
+        fallback=ctx.api.named_generic_type("builtins.function", []),
+    )
+
+    arg_types.append(calltyp)
+    arg_names.append("on_invalidate")
+    arg_kinds.append(ARG_NAMED_OPT)  # Arg is an optional kwarg.
+
+    signature = signature.copy_modified(
+        arg_types=arg_types, arg_names=arg_names, arg_kinds=arg_kinds,
+    )
 
     return signature
 
diff --git a/synapse/handlers/presence.py b/synapse/handlers/presence.py
index 49a00eed9c..8e014c9bb5 100644
--- a/synapse/handlers/presence.py
+++ b/synapse/handlers/presence.py
@@ -48,7 +48,7 @@ from synapse.util.wheel_timer import WheelTimer
 
 MYPY = False
 if MYPY:
-    import synapse.server
+    from synapse.server import HomeServer
 
 logger = logging.getLogger(__name__)
 
@@ -101,7 +101,7 @@ assert LAST_ACTIVE_GRANULARITY < IDLE_TIMER
 class BasePresenceHandler(abc.ABC):
     """Parts of the PresenceHandler that are shared between workers and master"""
 
-    def __init__(self, hs: "synapse.server.HomeServer"):
+    def __init__(self, hs: "HomeServer"):
         self.clock = hs.get_clock()
         self.store = hs.get_datastore()
 
@@ -199,7 +199,7 @@ class BasePresenceHandler(abc.ABC):
 
 
 class PresenceHandler(BasePresenceHandler):
-    def __init__(self, hs: "synapse.server.HomeServer"):
+    def __init__(self, hs: "HomeServer"):
         super().__init__(hs)
         self.hs = hs
         self.is_mine_id = hs.is_mine_id
@@ -1011,7 +1011,7 @@ def format_user_presence_state(state, now, include_user_id=True):
 
 
 class PresenceEventSource:
-    def __init__(self, hs):
+    def __init__(self, hs: "HomeServer"):
         # We can't call get_presence_handler here because there's a cycle:
         #
         #   Presence -> Notifier -> PresenceEventSource -> Presence
@@ -1071,12 +1071,14 @@ class PresenceEventSource:
 
             users_interested_in = await self._get_interested_in(user, explicit_room_id)
 
-            user_ids_changed = set()
+            user_ids_changed = set()  # type: Collection[str]
             changed = None
             if from_key:
                 changed = stream_change_cache.get_all_entities_changed(from_key)
 
             if changed is not None and len(changed) < 500:
+                assert isinstance(user_ids_changed, set)
+
                 # For small deltas, its quicker to get all changes and then
                 # work out if we share a room or they're in our presence list
                 get_updates_counter.labels("stream").inc()