Allow fancy self-types (#7860)

Fixes #3625
Fixes #5305
Fixes #5320
Fixes #5868
Fixes #7191
Fixes #7778
Fixes python/typing#680

So, lately I was noticing many issues that would be fixed by (partially) moving the check for self-type from definition site to call site.

This morning I found that we actually have such function `check_self_arg()` that is applied at call site, but it is almost not used. After more reading of the code I found that all the patterns for self-types that I wanted to support should either already work, or work with minimal modifications. Finally, I discovered that the root cause of many of the problems is the fact that `bind_self()` uses wrong direction for type inference! All these years it expected actual argument type to be _supertype_ of the formal one.

After fixing this bug, it turned out it was easy to support following patterns for explicit self-types:
* Structured match on generic self-types
* Restricted methods in generic classes (methods that one is allowed to call only for some values or type arguments)
* Methods overloaded on self-type
* (Important case of the above) overloaded `__init__` for generic classes
* Mixin classes (using protocols)
* Private class-level decorators (a bit hacky)
* Precise types for alternative constructors (mostly already worked)

This PR cuts few corners, but it is ready for review (I left some TODOs). Note I also add some docs, I am not sure this is really needed, but probably good to have.

Author: ilevkivskyi

docs/source/generics.rst

@@ -312,6 +312,8 @@ or a deserialization method returns the actual type of self. Therefore
 you may need to silence mypy inside these methods (but not at the call site),
 possibly by making use of the ``Any`` type.
 
+For some advanced uses of self-types see :ref:`additional examples <advanced_self>`.
+
 .. _variance-of-generics:
 
 Variance of generic types

docs/source/more_types.rst

@@ -561,6 +561,144 @@ with ``Union[int, slice]`` and ``Union[T, Sequence]``.
    to returning ``Any`` only if the input arguments also contain ``Any``.
 
 
+.. _advanced_self:
+
+Advanced uses of self-types
+***************************
+
+Normally, mypy doesn't require annotations for the first arguments of instance and
+class methods. However, they may be needed to have more precise static typing
+for certain programming patterns.
+
+Restricted methods in generic classes
+-------------------------------------
+
+In generic classes some methods may be allowed to be called only
+for certain values of type arguments:
+
+.. code-block:: python
+
+   T = TypeVar('T')
+
+   class Tag(Generic[T]):
+       item: T
+       def uppercase_item(self: C[str]) -> str:
+           return self.item.upper()
+
+   def label(ti: Tag[int], ts: Tag[str]) -> None:
+       ti.uppercase_item()  # E: Invalid self argument "Tag[int]" to attribute function
+                            # "uppercase_item" with type "Callable[[Tag[str]], str]"
+       ts.uppercase_item()  # This is OK
+
+This pattern also allows matching on nested types in situations where the type
+argument is itself generic:
+
+.. code-block:: python
+
+  T = TypeVar('T')
+  S = TypeVar('S')
+
+   class Storage(Generic[T]):
+       def __init__(self, content: T) -> None:
+           self.content = content
+       def first_chunk(self: Storage[Sequence[S]]) -> S:
+           return self.content[0]
+
+   page: Storage[List[str]]
+   page.first_chunk()  # OK, type is "str"
+
+   Storage(0).first_chunk()  # Error: Invalid self argument "Storage[int]" to attribute function
+                             # "first_chunk" with type "Callable[[Storage[Sequence[S]]], S]"
+
+Finally, one can use overloads on self-type to express precise types of
+some tricky methods:
+
+.. code-block:: python
+
+   T = TypeVar('T')
+
+   class Tag(Generic[T]):
+       @overload
+       def export(self: Tag[str]) -> str: ...
+       @overload
+       def export(self, converter: Callable[[T], str]) -> T: ...
+
+       def export(self, converter=None):
+           if isinstance(self.item, str):
+               return self.item
+           return converter(self.item)
+
+In particular, an :py:meth:`~object.__init__` method overloaded on self-type
+may be useful to annotate generic class constructors where type arguments
+depend on constructor parameters in a non-trivial way, see e.g. :py:class:`~subprocess.Popen`.
+
+Mixin classes
+-------------
+
+Using host class protocol as a self-type in mixin methods allows
+more code re-usability for static typing of mixin classes. For example,
+one can define a protocol that defines common functionality for
+host classes instead of adding required abstract methods to every mixin:
+
+.. code-block:: python
+
+   class Lockable(Protocol):
+       @property
+       def lock(self) -> Lock: ...
+
+   class AtomicCloseMixin:
+       def atomic_close(self: Lockable) -> int:
+           with self.lock:
+               # perform actions
+
+   class AtomicOpenMixin:
+       def atomic_open(self: Lockable) -> int:
+           with self.lock:
+               # perform actions
+
+   class File(AtomicCloseMixin, AtomicOpenMixin):
+       def __init__(self) -> None:
+           self.lock = Lock()
+
+   class Bad(AtomicCloseMixin):
+       pass
+
+   f = File()
+   b: Bad
+   f.atomic_close()  # OK
+   b.atomic_close()  # Error: Invalid self type for "atomic_close"
+
+Note that the explicit self-type is *required* to be a protocol whenever it
+is not a supertype of the current class. In this case mypy will check the validity
+of the self-type only at the call site.
+
+Precise typing of alternative constructors
+------------------------------------------
+
+Some classes may define alternative constructors. If these
+classes are generic, self-type allows giving them precise signatures:
+
+.. code-block:: python
+
+   T = TypeVar('T')
+
+   class Base(Generic[T]):
+       Q = TypeVar('Q', bound='Base[T]')
+
+       def __init__(self, item: T) -> None:
+           self.item = item
+
+       @classmethod
+       def make_pair(cls: Type[Q], item: T) -> Tuple[Q, Q]:
+           return cls(item), cls(item)
+
+   class Sub(Base[T]):
+       ...
+
+   pair = Sub.make_pair('yes')  # Type is "Tuple[Sub[str], Sub[str]]"
+   bad = Sub[int].make_pair('no')  # Error: Argument 1 to "make_pair" of "Base"
+                                   # has incompatible type "str"; expected "int"
+
 .. _async-and-await:
 
 Typing async/await

mypy/checker.py

@@ -906,10 +906,18 @@ def check_func_def(self, defn: FuncItem, typ: CallableType, name: Optional[str])
                         isclass = defn.is_class or defn.name() in ('__new__', '__init_subclass__')
                         if isclass:
                             ref_type = mypy.types.TypeType.make_normalized(ref_type)
-                        erased = erase_to_bound(arg_type)
+                        erased = get_proper_type(erase_to_bound(arg_type))
                         if not is_subtype_ignoring_tvars(ref_type, erased):
                             note = None
-                            if typ.arg_names[i] in ['self', 'cls']:
+                            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'}:
                                 if (self.options.python_version[0] < 3
                                         and is_same_type(erased, arg_type) and not isclass):
                                     msg = message_registry.INVALID_SELF_TYPE_OR_EXTRA_ARG
@@ -919,9 +927,10 @@ def check_func_def(self, defn: FuncItem, typ: CallableType, name: Optional[str])
                                         erased, ref_type)
                             else:
                                 msg = message_registry.MISSING_OR_INVALID_SELF_TYPE
-                            self.fail(msg, defn)
-                            if note:
-                                self.note(note, defn)
+                            if msg:
+                                self.fail(msg, defn)
+                                if note:
+                                    self.note(note, defn)
                     elif isinstance(arg_type, TypeVarType):
                         # Refuse covariant parameter type variables
                         # TODO: check recursively for inner type variables

mypy/checkmember.py

@@ -200,6 +200,12 @@ def analyze_instance_member_access(name: str,
             # the first argument.
             pass
         else:
+            if isinstance(signature, FunctionLike) and name != '__call__':
+                # TODO: use proper treatment of special methods on unions instead
+                #       of this hack here and below (i.e. mx.self_type).
+                dispatched_type = meet.meet_types(mx.original_type, typ)
+                signature = check_self_arg(signature, dispatched_type, False, mx.context,
+                                           name, mx.msg)
             signature = bind_self(signature, mx.self_type)
         typ = map_instance_to_supertype(typ, method.info)
         member_type = expand_type_by_instance(signature, typ)
@@ -546,8 +552,8 @@ def analyze_var(name: str,
                 # In `x.f`, when checking `x` against A1 we assume x is compatible with A
                 # and similarly for B1 when checking agains B
                 dispatched_type = meet.meet_types(mx.original_type, itype)
-                check_self_arg(functype, dispatched_type, var.is_classmethod, mx.context, name,
-                               mx.msg)
+                functype = check_self_arg(functype, dispatched_type, var.is_classmethod,
+                                          mx.context, name, mx.msg)
                 signature = bind_self(functype, mx.self_type, var.is_classmethod)
                 if var.is_property:
                     # A property cannot have an overloaded type => the cast is fine.
@@ -596,27 +602,45 @@ def check_self_arg(functype: FunctionLike,
                    dispatched_arg_type: Type,
                    is_classmethod: bool,
                    context: Context, name: str,
-                   msg: MessageBuilder) -> None:
-    """For x.f where A.f: A1 -> T, check that meet(type(x), A) <: A1 for each overload.
+                   msg: MessageBuilder) -> FunctionLike:
+    """Check that an instance has a valid type for a method with annotated 'self'.
 
-    dispatched_arg_type is meet(B, A) in the following example
-
-        def g(x: B): x.f
+    For example if the method is defined as:
         class A:
-            f: Callable[[A1], None]
+            def f(self: S) -> T: ...
+    then for 'x.f' we check that meet(type(x), A) <: S. If the method is overloaded, we
+    select only overloads items that satisfy this requirement. If there are no matching
+    overloads, an error is generated.
+
+    Note: dispatched_arg_type uses a meet to select a relevant item in case if the
+    original type of 'x' is a union. This is done because several special methods
+    treat union types in ad-hoc manner, so we can't use MemberContext.self_type yet.
     """
-    # TODO: this is too strict. We can return filtered overloads for matching definitions
-    for item in functype.items():
+    items = functype.items()
+    if not items:
+        return functype
+    new_items = []
+    for item in items:
         if not item.arg_types or item.arg_kinds[0] not in (ARG_POS, ARG_STAR):
             # No positional first (self) argument (*args is okay).
             msg.no_formal_self(name, item, context)
+            # This is pretty bad, so just return the original signature if
+            # there is at least one such error.
+            return functype
         else:
             selfarg = item.arg_types[0]
             if is_classmethod:
                 dispatched_arg_type = TypeType.make_normalized(dispatched_arg_type)
-            if not subtypes.is_subtype(dispatched_arg_type, erase_to_bound(selfarg)):
-                msg.incompatible_self_argument(name, dispatched_arg_type, item,
-                                               is_classmethod, context)
+            if subtypes.is_subtype(dispatched_arg_type, erase_typevars(erase_to_bound(selfarg))):
+                new_items.append(item)
+    if not new_items:
+        # Choose first item for the message (it may be not very helpful for overloads).
+        msg.incompatible_self_argument(name, dispatched_arg_type, items[0],
+                                       is_classmethod, context)
+        return functype
+    if len(new_items) == 1:
+        return new_items[0]
+    return Overloaded(new_items)
 
 
 def analyze_class_attribute_access(itype: Instance,
@@ -702,7 +726,10 @@ def analyze_class_attribute_access(itype: Instance,
 
         is_classmethod = ((is_decorated and cast(Decorator, node.node).func.is_class)
                           or (isinstance(node.node, FuncBase) and node.node.is_class))
-        result = add_class_tvars(get_proper_type(t), itype, isuper, is_classmethod,
+        t = get_proper_type(t)
+        if isinstance(t, FunctionLike) and is_classmethod:
+            t = check_self_arg(t, mx.self_type, False, mx.context, name, mx.msg)
+        result = add_class_tvars(t, itype, isuper, is_classmethod,
                                  mx.builtin_type, mx.self_type)
         if not mx.is_lvalue:
             result = analyze_descriptor_access(mx.original_type, result, mx.builtin_type,

mypy/infer.py

@@ -2,10 +2,11 @@
 
 from typing import List, Optional, Sequence
 
-from mypy.constraints import infer_constraints, infer_constraints_for_callable
+from mypy.constraints import (
+    infer_constraints, infer_constraints_for_callable, SUBTYPE_OF, SUPERTYPE_OF
+)
 from mypy.types import Type, TypeVarId, CallableType
 from mypy.solve import solve_constraints
-from mypy.constraints import SUBTYPE_OF
 
 
 def infer_function_type_arguments(callee_type: CallableType,
@@ -36,8 +37,10 @@ def infer_function_type_arguments(callee_type: CallableType,
 
 
 def infer_type_arguments(type_var_ids: List[TypeVarId],
-                         template: Type, actual: Type) -> List[Optional[Type]]:
+                         template: Type, actual: Type,
+                         is_supertype: bool = False) -> List[Optional[Type]]:
     # Like infer_function_type_arguments, but only match a single type
     # against a generic type.
-    constraints = infer_constraints(template, actual, SUBTYPE_OF)
+    constraints = infer_constraints(template, actual,
+                                    SUPERTYPE_OF if is_supertype else SUBTYPE_OF)
     return solve_constraints(type_var_ids, constraints)

mypy/meet.py

@@ -493,6 +493,10 @@ def visit_instance(self, t: Instance) -> ProperType:
             call = unpack_callback_protocol(t)
             if call:
                 return meet_types(call, self.s)
+        elif isinstance(self.s, FunctionLike) and self.s.is_type_obj() and t.type.is_metaclass():
+            if is_subtype(self.s.fallback, t):
+                return self.s
+            return self.default(self.s)
         elif isinstance(self.s, TypeType):
             return meet_types(t, self.s)
         elif isinstance(self.s, TupleType):