mypy plugin to check @cached return types

changelog.d/14911.misc

@@ -0,0 +1 @@
+Improve type hints.

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/federation/federation_server.py

@@ -710,7 +710,7 @@ async def on_send_join_request(
         state_event_ids: Collection[str]
         servers_in_room: Optional[Collection[str]]
         if caller_supports_partial_state:
-            summary = await self.store.get_room_summary(room_id)
+            summary = await self.store.get_room_summary(room_id)  # type: ignore[synapse-@cached-mutable]
             state_event_ids = _get_event_ids_for_partial_state_join(
                 event, prev_state_ids, summary
             )

synapse/handlers/relations.py

@@ -320,7 +320,7 @@ async def _get_threads_for_events(
         event_ids = [eid for eid in events_by_id.keys() if eid not in relations_by_id]
 
         # Fetch thread summaries.
-        summaries = await self._main_store.get_thread_summaries(event_ids)
+        summaries = await self._main_store.get_thread_summaries(event_ids)  # type: ignore[synapse-@cached-mutable]
 
         # Limit fetching whether the requester has participated in a thread to
         # events which are thread roots.
@@ -514,7 +514,7 @@ async def _fetch_edits() -> None:
             Note that there is no use in limiting edits by ignored users since the
             parent event should be ignored in the first place if the user is ignored.
             """
-            edits = await self._main_store.get_applicable_edits(
+            edits = await self._main_store.get_applicable_edits(  # type: ignore[synapse-@cached-mutable]
                 [
                     event_id
                     for event_id, event in events_by_id.items()

synapse/handlers/sync.py

@@ -811,7 +811,7 @@ async def compute_summary(
         )
 
         # this is heavily cached, thus: fast.
-        details = await self.store.get_room_summary(room_id)
+        details = await self.store.get_room_summary(room_id)  # type: ignore[synapse-@cached-mutable]
 
         name_id = state_ids.get((EventTypes.Name, ""))
         canonical_alias_id = state_ids.get((EventTypes.CanonicalAlias, ""))
@@ -1336,7 +1336,7 @@ async def unread_notifs_for_room_id(
             return RoomNotifCounts.empty()
 
         with Measure(self.clock, "unread_notifs_for_room_id"):
-            return await self.store.get_unread_event_push_actions_by_room_for_user(
+            return await self.store.get_unread_event_push_actions_by_room_for_user(  # type: ignore[synapse-@cached-mutable]
                 room_id,
                 sync_config.user.to_string(),
             )

synapse/storage/controllers/state.py

@@ -682,7 +682,7 @@ async def _get_joined_hosts(
         # `get_joined_hosts` is called with the "current" state group for the
         # room, and so consecutive calls will be for consecutive state groups
         # which point to the previous state group.
-        cache = await self.stores.main._get_joined_hosts_cache(room_id)
+        cache = await self.stores.main._get_joined_hosts_cache(room_id)  # type: ignore[synapse-@cached-mutable]
 
         # If the state group in the cache matches, we already have the data we need.
         if state_entry.state_group == cache.state_group:

synapse/storage/databases/main/event_push_actions.py

@@ -182,6 +182,7 @@ class UserPushAction(EmailPushAction):
     profile_tag: str
 
 
+# TODO This is used as a cached value and is mutable.
 @attr.s(slots=True, auto_attribs=True)
 class NotifCounts:
     """
@@ -193,15 +194,15 @@ class NotifCounts:
     highlight_count: int = 0
 
 
-@attr.s(slots=True, auto_attribs=True)
+@attr.s(slots=True, frozen=True, auto_attribs=True)
 class RoomNotifCounts:
     """
     The per-user, per-room count of notifications. Used by sync and push.
     """
 
     main_timeline: NotifCounts
     # Map of thread ID to the notification counts.
-    threads: Dict[str, NotifCounts]
+    threads: Mapping[str, NotifCounts]
 
     @staticmethod
     def empty() -> "RoomNotifCounts":

synapse/storage/databases/main/relations.py

@@ -516,6 +516,7 @@ def _get_references_for_events_txn(
     def get_applicable_edit(self, event_id: str) -> Optional[EventBase]:
         raise NotImplementedError()
 
+    # TODO: This returns a mutable object, which is generally bad.
     @cachedList(cached_method_name="get_applicable_edit", list_name="event_ids")
     async def get_applicable_edits(
         self, event_ids: Collection[str]
@@ -602,6 +603,7 @@ def _get_applicable_edits_txn(txn: LoggingTransaction) -> Dict[str, str]:
     def get_thread_summary(self, event_id: str) -> Optional[Tuple[int, EventBase]]:
         raise NotImplementedError()
 
+    # TODO: This returns a mutable object, which is generally bad.
     @cachedList(cached_method_name="get_thread_summary", list_name="event_ids")
     async def get_thread_summaries(
         self, event_ids: Collection[str]

synapse/storage/databases/main/roommember.py

@@ -1071,6 +1071,7 @@ async def _get_approximate_current_memberships_in_room(
         )
         return {row["event_id"]: row["membership"] for row in rows}
 
+    # TODO This returns a mutable object, which is generally confusing when using a cache.
     @cached(max_entries=10000)
     def _get_joined_hosts_cache(self, room_id: str) -> "_JoinedHostsCache":
         return _JoinedHostsCache()

synapse/storage/roommember.py

@@ -45,6 +45,7 @@ class ProfileInfo:
     display_name: Optional[str]
 
 
+# TODO This is used as a cached value and is mutable.
 @attr.s(slots=True, frozen=True, weakref_slot=False, auto_attribs=True)
 class MemberSummary:
     # A truncated list of (user_id, event_id) tuples for users of a given

tests/rest/admin/test_server_notice.py

@@ -438,7 +438,7 @@ def test_send_server_notice_delete_room(self) -> None:
 
         # It doesn't really matter what API we use here, we just want to assert
         # that the room doesn't exist.
-        summary = self.get_success(self.store.get_room_summary(first_room_id))
+        summary = self.get_success(self.store.get_room_summary(first_room_id))  # type: ignore[synapse-@cached-mutable]
         # The summary should be empty since the room doesn't exist.
         self.assertEqual(summary, {})
 

tests/rest/client/test_shadow_banned.py

@@ -182,7 +182,7 @@ def test_upgrade(self) -> None:
 
         # It doesn't really matter what API we use here, we just want to assert
         # that the room doesn't exist.
-        summary = self.get_success(self.store.get_room_summary(new_room_id))
+        summary = self.get_success(self.store.get_room_summary(new_room_id))  # type: ignore[synapse-@cached-mutable]
         # The summary should be empty since the room doesn't exist.
         self.assertEqual(summary, {})