mypy plugin to check @cached return types

scripts-dev/mypy_synapse_plugin.py

@@ -16,13 +16,24 @@
 can crop up, e.g the cache descriptors.
 """
 
-from typing import Callable, Optional, Type
+from typing import Callable, Optional, Tuple, Type
 
+import mypy.types
 from mypy.erasetype import remove_instance_last_known_values
-from mypy.nodes import ARG_NAMED_OPT
+from mypy.errorcodes import ErrorCode
+from mypy.nodes import ARG_NAMED_OPT, Var
 from mypy.plugin import MethodSigContext, Plugin
 from mypy.typeops import bind_self
-from mypy.types import CallableType, Instance, NoneType, UnionType
+from mypy.types import (
+    AnyType,
+    CallableType,
+    Instance,
+    NoneType,
+    TupleType,
+    TypeAliasType,
+    UninhabitedType,
+    UnionType,
+)
 
 
 class SynapsePlugin(Plugin):
@@ -49,10 +60,10 @@ def cached_function_method_signature(ctx: MethodSigContext) -> CallableType:
         3. an optional keyword argument `on_invalidated` should be added.
     """
 
-    # First we mark this as a bound function signature.
-    signature = bind_self(ctx.default_signature)
+    # 1. Mark this as a bound function signature.
+    signature: CallableType = bind_self(ctx.default_signature)
 
-    # Secondly, we remove any "cache_context" args.
+    # 2. Remove any "cache_context" args.
     #
     # Note: We should be only doing this if `cache_context=True` is set, but if
     # it isn't then the code will raise an exception when its called anyway, so
@@ -72,7 +83,7 @@ def cached_function_method_signature(ctx: MethodSigContext) -> CallableType:
         arg_names.pop(context_arg_index)
         arg_kinds.pop(context_arg_index)
 
-    # Third, we add an optional "on_invalidate" argument.
+    # 3. Add an optional "on_invalidate" argument.
     #
     # This is a either
     # - a callable which accepts no input and returns nothing, or
@@ -94,7 +105,7 @@ def cached_function_method_signature(ctx: MethodSigContext) -> CallableType:
     arg_names.append("on_invalidate")
     arg_kinds.append(ARG_NAMED_OPT)  # Arg is an optional kwarg.
 
-    # Finally we ensure the return type is a Deferred.
+    # 4. Ensure the return type is a Deferred.
     if (
         isinstance(signature.ret_type, Instance)
         and signature.ret_type.type.fullname == "twisted.internet.defer.Deferred"
@@ -131,9 +142,153 @@ def cached_function_method_signature(ctx: MethodSigContext) -> CallableType:
         ret_type=ret_type,
     )
 
+    # 5. Complain loudly if we are returning something mutable
+    check_is_cacheable(signature, ctx, ret_type)
+
     return signature
 
 
+def check_is_cacheable(
+    signature: CallableType,
+    ctx: MethodSigContext,
+    deferred_return_type: Instance,
+) -> None:
+    # The previous code wraps the return type into a Deferred.
+    assert deferred_return_type.type.fullname == "twisted.internet.defer.Deferred"
+    return_type = deferred_return_type.args[0]
+
+    verbose = ctx.api.options.verbosity >= 1
+    # TODO Technically a cachedList only needs immutable values, but forcing them
+    # to return Mapping instead of Dict is fine.
+    ok, note = is_cacheable(return_type, signature, verbose)
+
+    if ok:
+        message = f"function {signature.name} is @cached, returning {return_type}"
+    else:
+        message = f"function {signature.name} is @cached, but has mutable return value {return_type}"
+
+    if note:
+        message += f" ({note})"
+    message = message.replace("builtins.", "").replace("typing.", "")
+
+    # TODO The context is the context of the caller, not the method itself.
+    if ok and note:
+        ctx.api.note(message, ctx.context)  # type: ignore[attr-defined]
+    elif not ok:
+        ctx.api.fail(message, ctx.context, code=AT_CACHED_MUTABLE_RETURN)
+
+
+# Immutable simple values.
+IMMUTABLE_VALUE_TYPES = {
+    "builtins.bool",
+    "builtins.int",
+    "builtins.float",
+    "builtins.str",
+    "builtins.bytes",
+}
+
+# Types defined in Synapse which are known to be immutable.
+IMMUTABLE_CUSTOM_TYPES = {
+    "synapse.synapse_rust.push.FilteredPushRules",
+    # This is technically not immutable, but close enough.
+    "signedjson.types.VerifyKey",
+}
+
+# Immutable containers only if the values are also immutable.
+IMMUTABLE_CONTAINER_TYPES_REQUIRING_IMMUTABLE_ELEMENTS = {
+    "builtins.frozenset",
+    "builtins.tuple",
+    "typing.AbstractSet",
+    "typing.Sequence",
+    "immutabledict.immutabledict",
+}
+
+MUTABLE_CONTAINER_TYPES = {
+    "builtins.set",
+    "builtins.list",
+    "builtins.dict",
+}
+
+AT_CACHED_MUTABLE_RETURN = ErrorCode(
+    "synapse-@cached-mutable",
+    "@cached() should have an immutable return type",
+    "General",
+)
+
+
+def is_cacheable(
+    rt: mypy.types.Type, signature: CallableType, verbose: bool
+) -> Tuple[bool, Optional[str]]:
+    """
+    Returns: a 2-tuple (cacheable, message).
+        - cachable: False means the type is definitely not cacheable;
+            true means anything else.
+        - Optional message.
+    """
+
+    # This should probably be done via a TypeVisitor. Apologies to the reader!
+    if isinstance(rt, AnyType):
+        return True, ("may be mutable" if verbose else None)
+
+    elif isinstance(rt, Instance):
+        if (
+            rt.type.fullname in IMMUTABLE_VALUE_TYPES
+            or rt.type.fullname in IMMUTABLE_CUSTOM_TYPES
+        ):
+            return True, None
+
+        elif rt.type.fullname == "typing.Mapping":
+            return is_cacheable(rt.args[1], signature, verbose)
+
+        elif rt.type.fullname in IMMUTABLE_CONTAINER_TYPES_REQUIRING_IMMUTABLE_ELEMENTS:
+            # E.g. Collection[T] is cachable iff T is cachable.
+            return is_cacheable(rt.args[0], signature, verbose)
+
+        elif rt.type.fullname in MUTABLE_CONTAINER_TYPES:
+            return False, None
+
+        elif "attrs" in rt.type.metadata:
+            frozen = rt.type.metadata["attrs"]["frozen"]
+            if frozen:
+                for attribute in rt.type.metadata["attrs"]["attributes"]:
+                    attribute_name = attribute["name"]
+                    symbol_node = rt.type.names[attribute_name].node
+                    assert isinstance(symbol_node, Var)
+                    assert symbol_node.type is not None
+                    ok, note = is_cacheable(symbol_node.type, signature, verbose)
+                    if not ok:
+                        return False, f"non-frozen attrs property: {attribute_name}"
+                # All attributes were frozen.
+                return True, None
+            else:
+                return False, "non-frozen attrs class"
+
+        else:
+            return False, f"Don't know how to handle {rt.type.fullname}"
+
+    elif isinstance(rt, NoneType):
+        return True, None
+
+    elif isinstance(rt, (TupleType, UnionType)):
+        for item in rt.items:
+            ok, note = is_cacheable(item, signature, verbose)
+            if not ok:
+                return False, note
+        # This discards notes but that's probably fine
+        return True, None
+
+    elif isinstance(rt, TypeAliasType):
+        return is_cacheable(mypy.types.get_proper_type(rt), signature, verbose)
+
+    # The tests check what happens if you raise an Exception, so they don't return.
+    elif isinstance(rt, UninhabitedType) and rt.is_noreturn:
+        # There's no return value, just consider it cachable.
+        return True, None
+
+    else:
+        return False, f"Don't know how to handle {type(rt).__qualname__} return type"
+
+
 def plugin(version: str) -> Type[SynapsePlugin]:
     # This is the entry point of the plugin, and lets us deal with the fact
     # that the mypy plugin interface is *not* stable by looking at the version

synapse/api/filtering.py

@@ -37,7 +37,7 @@
 from synapse.api.errors import SynapseError
 from synapse.api.presence import UserPresenceState
 from synapse.events import EventBase, relation_from_event
-from synapse.types import JsonDict, RoomID, UserID
+from synapse.types import JsonDict, JsonMapping, RoomID, UserID
 
 if TYPE_CHECKING:
     from synapse.server import HomeServer
@@ -191,7 +191,7 @@ def check_valid_filter(self, user_filter_json: JsonDict) -> None:
 
 
 class FilterCollection:
-    def __init__(self, hs: "HomeServer", filter_json: JsonDict):
+    def __init__(self, hs: "HomeServer", filter_json: JsonMapping):
         self._filter_json = filter_json
 
         room_filter_json = self._filter_json.get("room", {})
@@ -219,7 +219,7 @@ def __init__(self, hs: "HomeServer", filter_json: JsonDict):
     def __repr__(self) -> str:
         return "<FilterCollection %s>" % (json.dumps(self._filter_json),)
 
-    def get_filter_json(self) -> JsonDict:
+    def get_filter_json(self) -> JsonMapping:
         return self._filter_json
 
     def timeline_limit(self) -> int:
@@ -313,7 +313,7 @@ def blocks_all_room_timeline(self) -> bool:
 
 
 class Filter:
-    def __init__(self, hs: "HomeServer", filter_json: JsonDict):
+    def __init__(self, hs: "HomeServer", filter_json: JsonMapping):
         self._hs = hs
         self._store = hs.get_datastores().main
         self.filter_json = filter_json

synapse/app/admin_cmd.py

@@ -17,7 +17,7 @@
 import os
 import sys
 import tempfile
-from typing import List, Mapping, Optional
+from typing import List, Mapping, Optional, Sequence
 
 from twisted.internet import defer, task
 
@@ -57,7 +57,7 @@
 from synapse.storage.databases.main.stream import StreamWorkerStore
 from synapse.storage.databases.main.tags import TagsWorkerStore
 from synapse.storage.databases.main.user_erasure_store import UserErasureWorkerStore
-from synapse.types import JsonDict, StateMap
+from synapse.types import JsonMapping, StateMap
 from synapse.util import SYNAPSE_VERSION
 from synapse.util.logcontext import LoggingContext
 
@@ -198,15 +198,15 @@ def write_knock(
             for event in state.values():
                 json.dump(event, fp=f)
 
-    def write_profile(self, profile: JsonDict) -> None:
+    def write_profile(self, profile: JsonMapping) -> None:
         user_directory = os.path.join(self.base_directory, "user_data")
         os.makedirs(user_directory, exist_ok=True)
         profile_file = os.path.join(user_directory, "profile")
 
         with open(profile_file, "a") as f:
             json.dump(profile, fp=f)
 
-    def write_devices(self, devices: List[JsonDict]) -> None:
+    def write_devices(self, devices: Sequence[JsonMapping]) -> None:
         user_directory = os.path.join(self.base_directory, "user_data")
         os.makedirs(user_directory, exist_ok=True)
         device_file = os.path.join(user_directory, "devices")
@@ -215,7 +215,7 @@ def write_devices(self, devices: List[JsonDict]) -> None:
             with open(device_file, "a") as f:
                 json.dump(device, fp=f)
 
-    def write_connections(self, connections: List[JsonDict]) -> None:
+    def write_connections(self, connections: Sequence[JsonMapping]) -> None:
         user_directory = os.path.join(self.base_directory, "user_data")
         os.makedirs(user_directory, exist_ok=True)
         connection_file = os.path.join(user_directory, "connections")
@@ -225,7 +225,7 @@ def write_connections(self, connections: List[JsonDict]) -> None:
                 json.dump(connection, fp=f)
 
     def write_account_data(
-        self, file_name: str, account_data: Mapping[str, JsonDict]
+        self, file_name: str, account_data: Mapping[str, JsonMapping]
     ) -> None:
         account_data_directory = os.path.join(
             self.base_directory, "user_data", "account_data"
@@ -237,7 +237,7 @@ def write_account_data(
         with open(account_data_file, "a") as f:
             json.dump(account_data, fp=f)
 
-    def write_media_id(self, media_id: str, media_metadata: JsonDict) -> None:
+    def write_media_id(self, media_id: str, media_metadata: JsonMapping) -> None:
         file_directory = os.path.join(self.base_directory, "media_ids")
         os.makedirs(file_directory, exist_ok=True)
         media_id_file = os.path.join(file_directory, media_id)

synapse/appservice/__init__.py

@@ -23,7 +23,7 @@
 
 from synapse.api.constants import EventTypes
 from synapse.events import EventBase
-from synapse.types import DeviceListUpdates, JsonDict, UserID
+from synapse.types import DeviceListUpdates, JsonDict, JsonMapping, UserID
 from synapse.util.caches.descriptors import _CacheContext, cached
 
 if TYPE_CHECKING:
@@ -379,8 +379,8 @@ def __init__(
         service: ApplicationService,
         id: int,
         events: Sequence[EventBase],
-        ephemeral: List[JsonDict],
-        to_device_messages: List[JsonDict],
+        ephemeral: List[JsonMapping],
+        to_device_messages: List[JsonMapping],
         one_time_keys_count: TransactionOneTimeKeysCount,
         unused_fallback_keys: TransactionUnusedFallbackKeys,
         device_list_summary: DeviceListUpdates,

synapse/appservice/api.py

@@ -41,7 +41,7 @@
 from synapse.events.utils import SerializeEventConfig, serialize_event
 from synapse.http.client import SimpleHttpClient, is_unknown_endpoint
 from synapse.logging import opentracing
-from synapse.types import DeviceListUpdates, JsonDict, ThirdPartyInstanceID
+from synapse.types import DeviceListUpdates, JsonDict, JsonMapping, ThirdPartyInstanceID
 from synapse.util.caches.response_cache import ResponseCache
 
 if TYPE_CHECKING:
@@ -306,8 +306,8 @@ async def push_bulk(
         self,
         service: "ApplicationService",
         events: Sequence[EventBase],
-        ephemeral: List[JsonDict],
-        to_device_messages: List[JsonDict],
+        ephemeral: List[JsonMapping],
+        to_device_messages: List[JsonMapping],
         one_time_keys_count: TransactionOneTimeKeysCount,
         unused_fallback_keys: TransactionUnusedFallbackKeys,
         device_list_summary: DeviceListUpdates,