Move self argument checks to a later phase (after decorator application, if any)

Author: sterliakov

mypy/checker.py

@@ -1369,49 +1369,19 @@ def check_func_def(
                         )
 
                 # Store argument types.
+                found_self = False
+                if isinstance(defn, FuncDef) and not defn.is_decorated:
+                    found_self = self.require_correct_self_argument(typ, defn)
                 for i in range(len(typ.arg_types)):
                     arg_type = typ.arg_types[i]
-                    if (
-                        isinstance(defn, FuncDef)
-                        and ref_type is not None
-                        and i == 0
-                        and defn.has_self_or_cls_argument
-                        and typ.arg_kinds[0] not in [nodes.ARG_STAR, nodes.ARG_STAR2]
-                    ):
-                        if defn.is_class or defn.name == "__new__":
-                            ref_type = mypy.types.TypeType.make_normalized(ref_type)
-                        if not is_same_type(arg_type, ref_type):
-                            # This level of erasure matches the one in checkmember.check_self_arg(),
-                            # better keep these two checks consistent.
-                            erased = get_proper_type(erase_typevars(erase_to_bound(arg_type)))
-                            if not is_subtype(ref_type, erased, ignore_type_params=True):
-                                if (
-                                    isinstance(erased, Instance)
-                                    and erased.type.is_protocol
-                                    or isinstance(erased, TypeType)
-                                    and isinstance(erased.item, Instance)
-                                    and erased.item.type.is_protocol
-                                ):
-                                    # We allow the explicit self-type to be not a supertype of
-                                    # the current class if it is a protocol. For such cases
-                                    # the consistency check will be performed at call sites.
-                                    msg = None
-                                elif typ.arg_names[i] in {"self", "cls"}:
-                                    msg = message_registry.ERASED_SELF_TYPE_NOT_SUPERTYPE.format(
-                                        erased.str_with_options(self.options),
-                                        ref_type.str_with_options(self.options),
-                                    )
-                                else:
-                                    msg = message_registry.MISSING_OR_INVALID_SELF_TYPE
-                                if msg:
-                                    self.fail(msg, defn)
-                    elif isinstance(arg_type, TypeVarType):
+                    if isinstance(arg_type, TypeVarType):
                         # Refuse covariant parameter type variables
                         # TODO: check recursively for inner type variables
                         if (
                             arg_type.variance == COVARIANT
                             and defn.name not in ("__init__", "__new__", "__post_init__")
                             and not is_private(defn.name)  # private methods are not inherited
+                            and (i != 0 or not found_self)
                         ):
                             ctx: Context = arg_type
                             if ctx.line < 0:
@@ -1561,6 +1531,60 @@ def check_func_def(
 
             self.binder = old_binder
 
+    def require_correct_self_argument(self, func: Type, defn: FuncDef) -> bool:
+        func = get_proper_type(func)
+        if not isinstance(func, CallableType):
+            return False
+
+        with self.scope.push_function(defn):
+            # We temporary push the definition to get the self type as
+            # visible from *inside* of this function/method.
+            ref_type: Type | None = self.scope.active_self_type()
+            if ref_type is None:
+                return False
+
+        if not defn.has_self_or_cls_argument or (
+            func.arg_kinds and func.arg_kinds[0] in [nodes.ARG_STAR, nodes.ARG_STAR2]
+        ):
+            return False
+
+        if not func.arg_types:
+            self.fail(
+                'Method must have at least one argument. Did you forget the "self" argument?', defn
+            )
+            return False
+
+        arg_type = func.arg_types[0]
+        if defn.is_class or defn.name == "__new__":
+            ref_type = mypy.types.TypeType.make_normalized(ref_type)
+        if is_same_type(arg_type, ref_type):
+            return True
+
+        # This level of erasure matches the one in checkmember.check_self_arg(),
+        # better keep these two checks consistent.
+        erased = get_proper_type(erase_typevars(erase_to_bound(arg_type)))
+        if not is_subtype(ref_type, erased, ignore_type_params=True):
+            if (
+                isinstance(erased, Instance)
+                and erased.type.is_protocol
+                or isinstance(erased, TypeType)
+                and isinstance(erased.item, Instance)
+                and erased.item.type.is_protocol
+            ):
+                # We allow the explicit self-type to be not a supertype of
+                # the current class if it is a protocol. For such cases
+                # the consistency check will be performed at call sites.
+                msg = None
+            elif func.arg_names[0] in {"self", "cls"}:
+                msg = message_registry.ERASED_SELF_TYPE_NOT_SUPERTYPE.format(
+                    erased.str_with_options(self.options), ref_type.str_with_options(self.options)
+                )
+            else:
+                msg = message_registry.MISSING_OR_INVALID_SELF_TYPE
+            if msg:
+                self.fail(msg, defn)
+        return True
+
     def is_var_redefined_in_outer_context(self, v: Var, after_line: int) -> bool:
         """Can the variable be assigned to at module top level or outer function?
 
@@ -5298,6 +5322,7 @@ def visit_decorator_inner(
             )
         if non_trivial_decorator:
             self.check_untyped_after_decorator(sig, e.func)
+        self.require_correct_self_argument(sig, e.func)
         sig = set_callable_name(sig, e.func)
         e.var.type = sig
         e.var.is_ready = True

mypy/semanal.py

@@ -1072,12 +1072,7 @@ def prepare_method_signature(self, func: FuncDef, info: TypeInfo, has_self_type:
         if func.has_self_or_cls_argument:
             if func.name in ["__init_subclass__", "__class_getitem__"]:
                 func.is_class = True
-            if not func.arguments:
-                self.fail(
-                    'Method must have at least one argument. Did you forget the "self" argument?',
-                    func,
-                )
-            elif isinstance(functype, CallableType):
+            if func.arguments and isinstance(functype, CallableType):
                 self_type = get_proper_type(functype.arg_types[0])
                 if isinstance(self_type, AnyType):
                     if has_self_type:

test-data/unit/check-functions.test

@@ -3708,3 +3708,62 @@ foo(*args)  # E: Argument 1 to "foo" has incompatible type "*list[object]"; expe
 kwargs: dict[str, object]
 foo(**kwargs)  # E: Argument 1 to "foo" has incompatible type "**dict[str, object]"; expected "P"
 [builtins fixtures/dict.pyi]
+
+[case testMethodDecoratorSelfChecks]
+from typing import Callable, ParamSpec, TypeVar
+
+T = TypeVar("T")
+P = ParamSpec("P")
+
+def to_number_1(fn: Callable[[], int]) -> int:
+    return 0
+
+def to_number_2(fn: Callable[[int], int]) -> int:
+    return 0
+
+def to_same_callable(fn: Callable[P, T]) -> Callable[P, T]:
+    return fn
+
+class A:
+    @to_number_1
+    def fn1() -> int:
+        return 0
+
+    @to_number_1  # E: Argument 1 to "to_number_1" has incompatible type "Callable[[int], int]"; expected "Callable[[], int]"
+    def fn2(_x: int) -> int:
+        return 0
+
+    @to_number_2  # E: Argument 1 to "to_number_2" has incompatible type "Callable[[], int]"; expected "Callable[[int], int]"
+    def fn3() -> int:
+        return 0
+
+    @to_number_2
+    def fn4(_x: int) -> int:
+        return 0
+
+    @to_number_2  # E: Argument 1 to "to_number_2" has incompatible type "Callable[[str], int]"; expected "Callable[[int], int]"
+    def fn5(_x: str) -> int:
+        return 0
+
+    @to_same_callable
+    def g1() -> None: ...  # E: Method must have at least one argument. Did you forget the "self" argument?
+
+    @to_same_callable
+    def g2(x: int) -> None: ...  # E: Self argument missing for a non-static method (or an invalid type for self)
+
+    @to_same_callable
+    def g3(self: int) -> None: ...  # E: The erased type of self "builtins.int" is not a supertype of its class "__main__.A"
+
+reveal_type(A().fn1)  # N: Revealed type is "builtins.int"
+reveal_type(A().fn2)  # N: Revealed type is "builtins.int"
+reveal_type(A().fn3)  # N: Revealed type is "builtins.int"
+reveal_type(A().fn4)  # N: Revealed type is "builtins.int"
+reveal_type(A().fn5)  # N: Revealed type is "builtins.int"
+
+reveal_type(A().g1)  # E: Attribute function "g1" with type "Callable[[], None]" does not accept self argument \
+                     # N: Revealed type is "def ()"
+reveal_type(A().g2)  # E: Invalid self argument "A" to attribute function "g2" with type "Callable[[int], None]" \
+                     # N: Revealed type is "def ()"
+reveal_type(A().g3)  # E: Invalid self argument "A" to attribute function "g3" with type "Callable[[int], None]" \
+                     # N: Revealed type is "def ()"
+[builtins fixtures/tuple.pyi]