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
|