✨ feat: add mypy plugin to enforce typedict return type on serializer methods

tools/mypy_helpers/plugin.py

@@ -4,9 +4,9 @@
 from collections.abc import Callable
 
 from mypy.build import PRI_MYPY
-from mypy.errorcodes import ATTR_DEFINED
+from mypy.errorcodes import ATTR_DEFINED, ErrorCode
 from mypy.messages import format_type
-from mypy.nodes import ARG_POS, MDEF, MypyFile, SymbolTableNode, TypeInfo, Var
+from mypy.nodes import ARG_POS, MDEF, FuncDef, MypyFile, SymbolTableNode, TypeInfo, Var
 from mypy.plugin import (
     AttributeContext,
     ClassDefContext,
@@ -24,10 +24,26 @@
     Instance,
     NoneType,
     Type,
+    TypedDictType,
     TypeOfAny,
     UnionType,
+    get_proper_type,
 )
 
+SERIALIZE_TYPEDDICT_CODE = ErrorCode(
+    "serialize-typeddict", "Serialize method must return a TypedDict", "attr-defined"
+)
+
+SERIALIZER_BASE_CLASS = "sentry.api.serializers.base.Serializer"
+
+
+def _is_typed_dict(typ: Type, api: SemanticAnalyzerPluginInterface) -> bool:
+    """Check if a type is a TypedDict"""
+    if isinstance(get_proper_type(typ), TypedDictType):
+        return True
+
+    return False
+
 
 def _make_using_required_str(ctx: FunctionSigContext) -> CallableType:
     sig = ctx.default_signature
@@ -173,6 +189,45 @@ def _lazy_service_wrapper_attribute(ctx: AttributeContext, *, attr: str) -> Type
         return member
 
 
+def _check_serializer_class(ctx: ClassDefContext) -> bool:
+    """
+    Hook that checks if subclasses of Serializer have a serialize method that returns a TypedDict.
+    """
+    if not ctx.cls.info.has_base(SERIALIZER_BASE_CLASS):
+        return True
+
+    # Look for the serialize method in the class
+    for name, node in ctx.cls.info.names.items():
+        if name == "serialize" and isinstance(node.node, FuncDef):
+            # First check if the method has a return type annotation at all
+            if not node.node.type:
+                ctx.api.fail(
+                    "Method 'serialize' must have an explicit return type annotation in classes inheriting from Serializer",
+                    node.node,
+                    code=SERIALIZE_TYPEDDICT_CODE,
+                )
+                continue
+
+            # Then check if it's a callable with a proper return type
+            if isinstance(node.node.type, CallableType):
+                ret_type = node.node.type.ret_type
+                if not _is_typed_dict(ret_type, ctx.api):
+                    ctx.api.fail(
+                        "Method 'serialize' must return a TypedDict in classes inheriting from Serializer",
+                        node.node,
+                        code=SERIALIZE_TYPEDDICT_CODE,
+                    )
+            else:
+                # This case should be rare, but handle it for completeness
+                ctx.api.fail(
+                    "Method 'serialize' has an invalid type annotation",
+                    node.node,
+                    code=SERIALIZE_TYPEDDICT_CODE,
+                )
+
+    return True
+
+
 class SentryMypyPlugin(Plugin):
     def get_function_signature_hook(
         self, fullname: str
@@ -190,6 +245,9 @@ def get_base_class_hook(self, fullname: str) -> Callable[[ClassDefContext], None
         else:
             return None
 
+    def get_class_decorator_hook_2(self, fullname: str) -> Callable[[ClassDefContext], bool] | None:
+        return _check_serializer_class
+
     def get_attribute_hook(self, fullname: str) -> Callable[[AttributeContext], Type] | None:
         if fullname.startswith("sentry.utils.lazy_service_wrapper.LazyServiceWrapper."):
             _, attr = fullname.rsplit(".", 1)