Allow fancy self-types

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,125 @@ 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 annotation for first argument 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 call 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 extraction of items in situations where type
+argument is itself generic:
+
+.. code-block:: python
+
+  T = TypeVar('T')
+  S = TypeVar('S')
+
+   class Node(Generic[T]):
+       def __init__(self, content: T) -> None:
+           self.content = content
+       def first_item(self: Node[Sequence[S]]) -> S:
+           return self.content[0]
+
+   page: Node[List[str]]
+   page.get_first_item()  # OK, type is "str"
+
+   Node(0).get_first_item()  # Error: Invalid self argument "Node[int]" to attribute function
+                             # "first_item" with type "Callable[[Node[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)
+
+Mixin classes
+-------------
+
+Using host class protocol as a self-type in mixin methods allows
+static typing of mixin class patter:
+
+.. code-block:: python
+
+   class Resource(Protocol):
+       def close(self) -> int: ...
+
+   class AtomicClose:
+       def atomic_close(self: Resource) -> int:
+           with Lock():
+               return self.close()
+
+   class File(AtomicClose):
+       def close(self) -> int:
+          ...
+
+   class Bad(AtomicClose):
+       ...
+
+   f: File
+   b: Bad
+   f.atomic_close()  # OK
+   b.atomic_close()  # Error: Invalid self type for "atomic_close"
+
+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')
+   Q = TypeVar('Q', bound=Base[Any])
+
+   class Base(Generic[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,12 @@ 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:
+                                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 +921,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,7 +602,7 @@ def check_self_arg(functype: FunctionLike,
                    dispatched_arg_type: Type,
                    is_classmethod: bool,
                    context: Context, name: str,
-                   msg: MessageBuilder) -> None:
+                   msg: MessageBuilder) -> FunctionLike:
     """For x.f where A.f: A1 -> T, check that meet(type(x), A) <: A1 for each overload.
 
     dispatched_arg_type is meet(B, A) in the following example
@@ -605,18 +611,31 @@ def g(x: B): x.f
         class A:
             f: Callable[[A1], None]
     """
-    # 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 +721,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,10 @@
 
 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
 from mypy.types import Type, TypeVarId, CallableType
 from mypy.solve import solve_constraints
-from mypy.constraints import SUBTYPE_OF
+from mypy.constraints import SUPERTYPE_OF
 
 
 def infer_function_type_arguments(callee_type: CallableType,
@@ -36,8 +36,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/typeops.py

@@ -9,13 +9,13 @@
 
 from mypy.types import (
     TupleType, Instance, FunctionLike, Type, CallableType, TypeVarDef, Overloaded,
-    TypeVarType, TypeType, UninhabitedType, FormalArgument, UnionType, NoneType,
+    TypeVarType, UninhabitedType, FormalArgument, UnionType, NoneType,
     AnyType, TypeOfAny, TypeType, ProperType, LiteralType, get_proper_type, get_proper_types,
     copy_type
 )
 from mypy.nodes import (
     FuncBase, FuncItem, OverloadedFuncDef, TypeInfo, TypeVar, ARG_STAR, ARG_STAR2, Expression,
-    StrExpr
+    StrExpr, ARG_POS
 )
 from mypy.maptype import map_instance_to_supertype
 from mypy.expandtype import expand_type_by_instance, expand_type
@@ -50,6 +50,15 @@ def type_object_type_from_function(signature: FunctionLike,
     #   class B(A[List[T]], Generic[T]): pass
     #
     # We need to first map B's __init__ to the type (List[T]) -> None.
+
+    # We first record all non-trivial (explicit) self types.
+    if not is_new and def_info == info and not info.is_newtype:
+        orig_self_types = [(it.arg_types[0] if it.arg_types and it.arg_kinds[0] == ARG_POS and
+                            it.arg_types[0] != fill_typevars(def_info) else None)
+                           for it in signature.items()]
+    else:
+        orig_self_types = [None] * len(signature.items())
+
     signature = bind_self(signature, original_type=fill_typevars(info), is_classmethod=is_new)
     signature = cast(FunctionLike,
                      map_type_from_supertype(signature, info, def_info))
@@ -59,19 +68,19 @@ def type_object_type_from_function(signature: FunctionLike,
         special_sig = 'dict'
 
     if isinstance(signature, CallableType):
-        return class_callable(signature, info, fallback, special_sig, is_new)
+        return class_callable(signature, info, fallback, special_sig, is_new, orig_self_types[0])
     else:
         # Overloaded __init__/__new__.
         assert isinstance(signature, Overloaded)
         items = []  # type: List[CallableType]
-        for item in signature.items():
-            items.append(class_callable(item, info, fallback, special_sig, is_new))
+        for item, orig_self in zip(signature.items(), orig_self_types):
+            items.append(class_callable(item, info, fallback, special_sig, is_new, orig_self))
         return Overloaded(items)
 
 
 def class_callable(init_type: CallableType, info: TypeInfo, type_type: Instance,
                    special_sig: Optional[str],
-                   is_new: bool) -> CallableType:
+                   is_new: bool, orig_self_type: Optional[Type] = None) -> CallableType:
     """Create a type object type based on the signature of __init__."""
     variables = []  # type: List[TypeVarDef]
     variables.extend(info.defn.type_vars)
@@ -89,6 +98,8 @@ def class_callable(init_type: CallableType, info: TypeInfo, type_type: Instance,
         and is_subtype(init_ret_type, default_ret_type, ignore_type_params=True)
     ):
         ret_type = init_ret_type  # type: Type
+    elif orig_self_type is not None:
+        ret_type = orig_self_type
     else:
         ret_type = default_ret_type
 
@@ -130,6 +141,12 @@ def map_type_from_supertype(typ: Type,
     return expand_type_by_instance(typ, inst_type)
 
 
+def instance_or_var(typ: ProperType) -> bool:
+    # TODO: use more principled check for non-trivial self-types.
+    return (isinstance(typ, TypeVarType) or
+            isinstance(typ, Instance) and typ != fill_typevars(typ.type))
+
+
 F = TypeVar('F', bound=FunctionLike)
 
 
@@ -174,23 +191,25 @@ class B(A): pass
         # TODO: infer bounds on the type of *args?
         return cast(F, func)
     self_param_type = get_proper_type(func.arg_types[0])
-    if func.variables and (isinstance(self_param_type, TypeVarType) or
+    if func.variables and (instance_or_var(self_param_type) or
                            (isinstance(self_param_type, TypeType) and
-                            isinstance(self_param_type.item, TypeVarType))):
+                            instance_or_var(self_param_type.item))):
         if original_type is None:
             # Type check method override
             # XXX value restriction as union?
             original_type = erase_to_bound(self_param_type)
         original_type = get_proper_type(original_type)
 
         ids = [x.id for x in func.variables]
-        typearg = get_proper_type(infer_type_arguments(ids, self_param_type, original_type)[0])
+        typearg = get_proper_type(infer_type_arguments(ids, self_param_type,
+                                                       original_type, is_supertype=True)[0])
         if (is_classmethod and isinstance(typearg, UninhabitedType)
                 and isinstance(original_type, (Instance, TypeVarType, TupleType))):
             # In case we call a classmethod through an instance x, fallback to type(x)
             # TODO: handle Union
             typearg = get_proper_type(infer_type_arguments(ids, self_param_type,
-                                                           TypeType(original_type))[0])
+                                                           TypeType(original_type),
+                                                           is_supertype=True)[0])
 
         def expand(target: Type) -> Type:
             assert typearg is not None

test-data/unit/check-classes.test

@@ -4286,9 +4286,9 @@ reveal_type(ta.g3)  # N: Revealed type is 'def () -> Type[__main__.A]'
 reveal_type(ta.g4)  # N: Revealed type is 'def () -> Type[__main__.A]'
 
 x: M = ta
-x.g1  # should be error: Argument 0 to "g1" of "M" has incompatible type "M"; expected "Type[A]"
-x.g2  # should be error: Argument 0 to "g2" of "M" has incompatible type "M"; expected "Type[TA]"
-x.g3  # should be error: Argument 0 to "g3" of "M" has incompatible type "M"; expected "TTA"
+x.g1  # E: Invalid self argument "M" to attribute function "g1" with type "Callable[[Type[A]], A]"
+x.g2  # E: Invalid self argument "M" to attribute function "g2" with type "Callable[[Type[TA]], TA]"
+x.g3  # E: Invalid self argument "M" to attribute function "g3" with type "Callable[[TTA], TTA]"
 reveal_type(x.g4)  # N: Revealed type is 'def () -> __main__.M*'
 
 def r(ta: Type[TA], tta: TTA) -> None: