ref: future-proof using signature modification (#53347)

the signature of a few of these changes in newer `django-stubs` -- this
makes the plugin a little more future-proof

Author: asottile-sentry

tools/mypy_helpers/plugin.py

@@ -7,71 +7,52 @@
 from mypy.types import CallableType, FunctionLike, Instance
 
 
-def replace_transaction_atomic_sig_callback(ctx: FunctionSigContext) -> CallableType:
-    signature = ctx.default_signature
+def _make_using_required_str(ctx: FunctionSigContext) -> CallableType:
+    sig = ctx.default_signature
 
-    using_arg = signature.argument_by_name("using")
-    if not using_arg:
-        # No using arg in the signature, bail
-        return signature
+    using_arg = sig.argument_by_name("using")
+    if using_arg is None or using_arg.pos is None:
+        ctx.api.fail("The using parameter is required", ctx.context)
+        return sig
 
-    # We care about context managers.
-    ret_type = signature.ret_type
-    if not isinstance(ret_type, Instance):
-        return signature
+    for kind in sig.arg_kinds[: using_arg.pos]:
+        if kind != ARG_POS:
+            ctx.api.fail("Expected using to be the first optional", ctx.context)
+            return sig
 
-    # Replace the type and remove the default value of using.
     str_type = ctx.api.named_generic_type("builtins.str", [])
+    arg_kinds = [*sig.arg_kinds[: using_arg.pos], ARG_POS, *sig.arg_kinds[using_arg.pos + 1 :]]
+    arg_types = [*sig.arg_types[: using_arg.pos], str_type, *sig.arg_types[using_arg.pos + 1 :]]
+    return sig.copy_modified(arg_kinds=arg_kinds, arg_types=arg_types)
 
-    arg_types = signature.arg_types[1:]
-    arg_kinds = signature.arg_kinds[1:]
-
-    return signature.copy_modified(
-        arg_kinds=[ARG_POS, *arg_kinds],
-        arg_types=[str_type, *arg_types],
-    )
-
-
-def replace_get_connection_sig_callback(ctx: FunctionSigContext) -> CallableType:
-    signature = ctx.default_signature
-    using_arg = signature.argument_by_name("using")
-    if not using_arg:
-        ctx.api.fail("The using parameter is required", ctx.context)
 
-    str_type = ctx.api.named_generic_type("builtins.str", [])
+def replace_transaction_atomic_sig_callback(ctx: FunctionSigContext) -> CallableType:
+    sig = ctx.default_signature
 
-    return signature.copy_modified(arg_kinds=[ARG_POS], arg_types=[str_type])
+    if not sig.argument_by_name("using"):
+        # No using arg in the signature, bail
+        return sig
 
+    # We care about context managers.
+    if not isinstance(sig.ret_type, Instance):
+        return sig
 
-def replace_trailing_using_sig_callback(ctx: FunctionSigContext) -> CallableType:
-    signature = ctx.default_signature
-    using_arg = signature.argument_by_name("using")
-    if not using_arg:
-        ctx.api.fail("The using parameter is required", ctx.context)
+    return _make_using_required_str(ctx)
 
-    # Update the parameter type to be required and str
-    str_type = ctx.api.named_generic_type("builtins.str", [])
-    arg_kinds = signature.arg_kinds[0:-1]
-    arg_types = signature.arg_types[0:-1]
 
-    return signature.copy_modified(
-        arg_kinds=[*arg_kinds, ARG_POS], arg_types=[*arg_types, str_type]
-    )
+_FUNCTION_SIGNATURE_HOOKS = {
+    "django.db.transaction.atomic": replace_transaction_atomic_sig_callback,
+    "django.db.transaction.get_connection": _make_using_required_str,
+    "django.db.transaction.on_commit": _make_using_required_str,
+    "django.db.transaction.set_rollback": _make_using_required_str,
+}
 
 
 class SentryMypyPlugin(Plugin):
     def get_function_signature_hook(
         self, fullname: str
     ) -> Callable[[FunctionSigContext], FunctionLike] | None:
-        if fullname == "django.db.transaction.atomic":
-            return replace_transaction_atomic_sig_callback
-        if fullname == "django.db.transaction.get_connection":
-            return replace_get_connection_sig_callback
-        if fullname == "django.db.transaction.on_commit":
-            return replace_trailing_using_sig_callback
-        if fullname == "django.db.transaction.set_rollback":
-            return replace_trailing_using_sig_callback
-        return None
+        return _FUNCTION_SIGNATURE_HOOKS.get(fullname)
 
 
 def plugin(version: str) -> type[SentryMypyPlugin]: