Just expand self type together with the callable

Author: sterliakov

mypy/checkexpr.py

@@ -2916,7 +2916,6 @@ def infer_overload_return_type(
         # If we have a selftype overload, it should contribute to `any_causes_overload_ambiguity`
         # check. Pretend that we're checking `Foo.func(instance, ...)` instead of
         # `instance.func(...)`.
-        p_object_type = get_proper_type(object_type) if object_type is not None else None
 
         def is_trivial_self(t: CallableType) -> bool:
             if isinstance(t.definition, FuncDef):
@@ -2925,13 +2924,9 @@ def is_trivial_self(t: CallableType) -> bool:
                 return t.definition.func.is_trivial_self
             return False
 
-        prepend_self = (
-            isinstance(p_object_type, Instance)
-            and has_any_type(p_object_type)
-            and any(
-                typ.is_bound and typ.original_self_type is not None and not is_trivial_self(typ)
-                for typ in plausible_targets
-            )
+        prepend_self = has_any_type(object_type) and any(
+            typ.is_bound and typ.original_self_type is not None and not is_trivial_self(typ)
+            for typ in plausible_targets
         )
         if prepend_self:
             assert object_type is not None
@@ -2952,16 +2947,15 @@ def maybe_bind_self(t: Type) -> Type:
                 if prepend_self:
                     param = typ.original_self_type
                     assert param is not None, "Overload bound only partially?"
-                    assert isinstance(p_object_type, Instance)
-                    param = expand_type_by_instance(param, p_object_type)
+                    typ = typ.copy_modified(
+                        arg_types=[param] + typ.arg_types,
+                        arg_kinds=[ARG_POS] + typ.arg_kinds,
+                        arg_names=[None, *typ.arg_names],
+                        is_bound=False,
+                        original_self_type=None,
+                    )
                     ret_type, infer_type = self.check_call(
-                        callee=typ.copy_modified(
-                            arg_types=[param] + typ.arg_types,
-                            arg_kinds=[ARG_POS] + typ.arg_kinds,
-                            arg_names=[None, *typ.arg_names],
-                            is_bound=False,
-                            original_self_type=None,
-                        ),
+                        callee=typ,
                         args=args,
                         arg_kinds=arg_kinds,
                         arg_names=arg_names,

mypy/expandtype.py

@@ -381,6 +381,11 @@ def visit_callable_type(self, t: CallableType) -> CallableType:
                     arg_kinds=t.arg_kinds[:-2] + repl.arg_kinds,
                     arg_names=t.arg_names[:-2] + repl.arg_names,
                     ret_type=t.ret_type.accept(self),
+                    original_self_type=(
+                        t.original_self_type.accept(self)
+                        if t.original_self_type is not None
+                        else None
+                    ),
                     type_guard=(t.type_guard.accept(self) if t.type_guard is not None else None),
                     type_is=(t.type_is.accept(self) if t.type_is is not None else None),
                     imprecise_arg_kinds=(t.imprecise_arg_kinds or repl.imprecise_arg_kinds),
@@ -420,6 +425,9 @@ def visit_callable_type(self, t: CallableType) -> CallableType:
         expanded = t.copy_modified(
             arg_types=arg_types,
             ret_type=t.ret_type.accept(self),
+            original_self_type=(
+                t.original_self_type.accept(self) if t.original_self_type is not None else None
+            ),
             type_guard=(t.type_guard.accept(self) if t.type_guard is not None else None),
             type_is=(t.type_is.accept(self) if t.type_is is not None else None),
         )

mypy/typetraverser.py

@@ -87,9 +87,11 @@ def visit_callable_type(self, t: CallableType, /) -> None:
         t.ret_type.accept(self)
         t.fallback.accept(self)
 
+        if t.original_self_type is not None:
+            t.original_self_type.accept(self)
+
         if t.type_guard is not None:
             t.type_guard.accept(self)
-
         if t.type_is is not None:
             t.type_is.accept(self)