Implement generic type aliases (#2378)

Fixes #606 as per PEP 484.

Now type aliases may be generic, so that the example in PEP works. Generic type aliases are allowed for generic classes, unions, callables, and tuples. For example:

Vec = Iterable[Tuple[T, T]]
TInt = Tuple[T, int]
UInt = Union[T, int]
CBack = Callable[..., T]

The aliases could be used as specified in PEP 484, e.g. either one specifies all free type variables, or if unspecified they are all substituted by Any, for example:

def fun(v: Vec[T]) -> Vec[T]: # Same as Iterable[Tuple[T, T]]
    ...
v1: Vec[int] = []      # Same as Iterable[Tuple[int, int]]
v2: Vec = []           # Same as Iterable[Tuple[Any, Any]]
v3: Vec[int, int] = [] # Error: Invalid alias, too many type arguments!

Generic aliases may be used everywhere where a normal generic type could be used (in annotations, generic base classes etc, and as of recently also in runtime expressions). Nested and chained aliases are allowed, but excessive use of those is discouraged in the docs. Like ordinary (non-generic) aliases, generic ones may be imported from other modules.

NOTE: Many examples in the tests and in docs require a recent version of typing.py from python/typing to work at runtime (see #2382). This feature may be used with older versions of typing.py by using type comments or "forward references".

Author: ilevkivskyi

docs/source/kinds_of_types.rst

@@ -426,9 +426,69 @@ assigning the type to a variable:
    def f() -> AliasType:
        ...
 
-A type alias does not create a new type. It's just a shorthand notation
-for another type -- it's equivalent to the target type. Type aliases
-can be imported from modules like any names.
+Type aliases can be generic, in this case they could be used in two variants:
+Subscripted aliases are equivalent to original types with substituted type variables,
+number of type arguments must match the number of free type variables
+in generic type alias. Unsubscripted aliases are treated as original types with free
+variables replaced with ``Any``. Examples (following `PEP 484
+<https://www.python.org/dev/peps/pep-0484/#type-aliases>`_):
+
+.. code-block:: python
+
+    from typing import TypeVar, Iterable, Tuple, Union, Callable
+    S = TypeVar('S')
+    TInt = Tuple[int, S]
+    UInt = Union[S, int]
+    CBack = Callable[..., S]
+
+    def response(query: str) -> UInt[str]:  # Same as Union[str, int]
+        ...
+    def activate(cb: CBack[S]) -> S:        # Same as Callable[..., S]
+        ...
+    table_entry: TInt  # Same as Tuple[int, Any]
+
+    T = TypeVar('T', int, float, complex)
+    Vec = Iterable[Tuple[T, T]]
+
+    def inproduct(v: Vec[T]) -> T:
+        return sum(x*y for x, y in v)
+
+    def dilate(v: Vec[T], scale: T) -> Vec[T]:
+        return ((x * scale, y * scale) for x, y in v)
+
+    v1: Vec[int] = []      # Same as Iterable[Tuple[int, int]]
+    v2: Vec = []           # Same as Iterable[Tuple[Any, Any]]
+    v3: Vec[int, int] = [] # Error: Invalid alias, too many type arguments!
+
+Type aliases can be imported from modules like any names. Aliases can target another
+aliases (although building complex chains of aliases is not recommended, this
+impedes code readability, thus defeating the purpose of using aliases).
+Following previous examples:
+
+.. code-block:: python
+
+    from typing import TypeVar, Generic, Optional
+    from first_example import AliasType
+    from second_example import Vec
+
+    def fun() -> AliasType:
+        ...
+
+    T = TypeVar('T')
+    class NewVec(Generic[T], Vec[T]):
+        ...
+    for i, j in NewVec[int]():
+        ...
+
+    OIntVec = Optional[Vec[int]]
+
+.. note::
+
+    A type alias does not create a new type. It's just a shorthand notation for
+    another type -- it's equivalent to the target type. For generic type aliases
+    this means that variance of type variables used for alias definition does not
+    apply to aliases. A parameterized generic alias is treated simply as an original
+    type with the corresponding type variables substituted.
 
 .. _newtypes:
 

mypy/checkexpr.py

@@ -6,7 +6,8 @@
     Type, AnyType, CallableType, Overloaded, NoneTyp, Void, TypeVarDef,
     TupleType, Instance, TypeVarId, TypeVarType, ErasedType, UnionType,
     PartialType, DeletedType, UnboundType, UninhabitedType, TypeType,
-    true_only, false_only, is_named_instance, function_type
+    true_only, false_only, is_named_instance, function_type,
+    get_typ_args, set_typ_args,
 )
 from mypy.nodes import (
     NameExpr, RefExpr, Var, FuncDef, OverloadedFuncDef, TypeInfo, CallExpr,
@@ -17,6 +18,7 @@
     ConditionalExpr, ComparisonExpr, TempNode, SetComprehension,
     DictionaryComprehension, ComplexExpr, EllipsisExpr, StarExpr,
     TypeAliasExpr, BackquoteExpr, ARG_POS, ARG_NAMED, ARG_STAR2, MODULE_REF,
+    UNBOUND_TVAR, BOUND_TVAR,
 )
 from mypy import nodes
 import mypy.checker
@@ -1375,8 +1377,55 @@ def visit_type_application(self, tapp: TypeApplication) -> Type:
         return AnyType()
 
     def visit_type_alias_expr(self, alias: TypeAliasExpr) -> Type:
+        """Get type of a type alias (could be generic) in a runtime expression."""
+        item = alias.type
+        if not alias.in_runtime:
+            # We don't replace TypeVar's with Any for alias used as Alias[T](42).
+            item = self.replace_tvars_any(item)
+        if isinstance(item, Instance):
+            # Normally we get a callable type (or overloaded) with .is_type_obj() true
+            # representing the class's constructor
+            tp = type_object_type(item.type, self.named_type)
+        else:
+            # This type is invalid in most runtime contexts
+            # and corresponding an error will be reported.
+            return alias.fallback
+        if isinstance(tp, CallableType):
+            if len(tp.variables) != len(item.args):
+                self.msg.incompatible_type_application(len(tp.variables),
+                                                       len(item.args), item)
+                return AnyType()
+            return self.apply_generic_arguments(tp, item.args, item)
+        elif isinstance(tp, Overloaded):
+            for it in tp.items():
+                if len(it.variables) != len(item.args):
+                    self.msg.incompatible_type_application(len(it.variables),
+                                                           len(item.args), item)
+                    return AnyType()
+            return Overloaded([self.apply_generic_arguments(it, item.args, item)
+                               for it in tp.items()])
         return AnyType()
 
+    def replace_tvars_any(self, tp: Type) -> Type:
+        """Replace all type variables of a type alias tp with Any. Basically, this function
+        finishes what could not be done in method TypeAnalyser.visit_unbound_type()
+        from typeanal.py.
+        """
+        typ_args = get_typ_args(tp)
+        new_args = typ_args[:]
+        for i, arg in enumerate(typ_args):
+            if isinstance(arg, UnboundType):
+                sym = None
+                try:
+                    sym = self.chk.lookup_qualified(arg.name)
+                except KeyError:
+                    pass
+                if sym and (sym.kind == UNBOUND_TVAR or sym.kind == BOUND_TVAR):
+                    new_args[i] = AnyType()
+            else:
+                new_args[i] = self.replace_tvars_any(arg)
+        return set_typ_args(tp, new_args, tp.line, tp.column)
+
     def visit_list_expr(self, e: ListExpr) -> Type:
         """Type check a list expression [...]."""
         return self.check_lst_expr(e.items, 'builtins.list', '<list>', e)

mypy/exprtotype.py

@@ -20,11 +20,11 @@ def expr_to_unanalyzed_type(expr: Expression) -> Type:
     """
     if isinstance(expr, NameExpr):
         name = expr.name
-        return UnboundType(name, line=expr.line)
+        return UnboundType(name, line=expr.line, column=expr.column)
     elif isinstance(expr, MemberExpr):
         fullname = get_member_expr_fullname(expr)
         if fullname:
-            return UnboundType(fullname, line=expr.line)
+            return UnboundType(fullname, line=expr.line, column=expr.column)
         else:
             raise TypeTranslationError()
     elif isinstance(expr, IndexExpr):
@@ -42,7 +42,7 @@ def expr_to_unanalyzed_type(expr: Expression) -> Type:
             raise TypeTranslationError()
     elif isinstance(expr, ListExpr):
         return TypeList([expr_to_unanalyzed_type(t) for t in expr.items],
-                        line=expr.line)
+                        line=expr.line, column=expr.column)
     elif isinstance(expr, (StrExpr, BytesExpr)):
         # Parse string literal type.
         try:

mypy/fastparse.py

@@ -432,6 +432,7 @@ def visit_Assign(self, n: ast35.Assign) -> AssignmentStmt:
             typ = parse_type_comment(n.type_comment, n.lineno)
         elif new_syntax:
             typ = TypeConverter(line=n.lineno).visit(n.annotation)  # type: ignore
+            typ.column = n.annotation.col_offset
         if n.value is None:  # always allow 'x: int'
             rvalue = TempNode(AnyType())  # type: Expression
         else:

mypy/nodes.py

@@ -1732,9 +1732,18 @@ class TypeAliasExpr(Expression):
     """Type alias expression (rvalue)."""
 
     type = None  # type: mypy.types.Type
-
-    def __init__(self, type: 'mypy.types.Type') -> None:
+    # Simple fallback type for aliases that are invalid in runtime expressions
+    # (for example Union, Tuple, Callable).
+    fallback = None  # type: mypy.types.Type
+    # This type alias is subscripted in a runtime expression like Alias[int](42)
+    # (not in a type context like type annotation or base class).
+    in_runtime = False  # type: bool
+
+    def __init__(self, type: 'mypy.types.Type', fallback: 'mypy.types.Type' = None,
+                 in_runtime: bool = False) -> None:
         self.type = type
+        self.fallback = fallback
+        self.in_runtime = in_runtime
 
     def accept(self, visitor: NodeVisitor[T]) -> T:
         return visitor.visit_type_alias_expr(self)

mypy/semanal.py

@@ -1126,7 +1126,8 @@ def visit_block_maybe(self, b: Block) -> None:
         if b:
             self.visit_block(b)
 
-    def anal_type(self, t: Type, allow_tuple_literal: bool = False) -> Type:
+    def anal_type(self, t: Type, allow_tuple_literal: bool = False,
+                  aliasing: bool = False) -> Type:
         if t:
             if allow_tuple_literal:
                 # Types such as (t1, t2, ...) only allowed in assignment statements. They'll
@@ -1143,7 +1144,8 @@ def anal_type(self, t: Type, allow_tuple_literal: bool = False) -> Type:
                     return TupleType(items, self.builtin_type('builtins.tuple'), t.line)
             a = TypeAnalyser(self.lookup_qualified,
                              self.lookup_fully_qualified,
-                             self.fail)
+                             self.fail,
+                             aliasing=aliasing)
             return t.accept(a)
         else:
             return None
@@ -1173,7 +1175,8 @@ def visit_assignment_stmt(self, s: AssignmentStmt) -> None:
                     node.kind = TYPE_ALIAS
                     node.type_override = res
                     if isinstance(s.rvalue, IndexExpr):
-                        s.rvalue.analyzed = TypeAliasExpr(res)
+                        s.rvalue.analyzed = TypeAliasExpr(res,
+                                                          fallback=self.alias_fallback(res))
         if s.type:
             # Store type into nodes.
             for lvalue in s.lvalues:
@@ -1211,6 +1214,19 @@ def analyze_simple_literal_type(self, rvalue: Expression) -> Optional[Type]:
             return self.named_type_or_none('builtins.unicode')
         return None
 
+    def alias_fallback(self, tp: Type) -> Instance:
+        """Make a dummy Instance with no methods. It is used as a fallback type
+        to detect errors for non-Instance aliases (i.e. Unions, Tuples, Callables).
+        """
+        kind = (' to Callable' if isinstance(tp, CallableType) else
+                ' to Tuple' if isinstance(tp, TupleType) else
+                ' to Union' if isinstance(tp, UnionType) else '')
+        cdef = ClassDef('Type alias' + kind, Block([]))
+        fb_info = TypeInfo(SymbolTable(), cdef, self.cur_mod_id)
+        fb_info.bases = [self.object_type()]
+        fb_info.mro = [fb_info, self.object_type().type]
+        return Instance(fb_info, [])
+
     def check_and_set_up_type_alias(self, s: AssignmentStmt) -> None:
         """Check if assignment creates a type alias and set it up as needed."""
         # For now, type aliases only work at the top level of a module.
@@ -2361,7 +2377,16 @@ def visit_unary_expr(self, expr: UnaryExpr) -> None:
 
     def visit_index_expr(self, expr: IndexExpr) -> None:
         expr.base.accept(self)
-        if refers_to_class_or_function(expr.base):
+        if isinstance(expr.base, RefExpr) and expr.base.kind == TYPE_ALIAS:
+            # Special form -- subscripting a generic type alias.
+            # Perform the type substitution and create a new alias.
+            res = analyze_type_alias(expr,
+                                     self.lookup_qualified,
+                                     self.lookup_fully_qualified,
+                                     self.fail)
+            expr.analyzed = TypeAliasExpr(res, fallback=self.alias_fallback(res),
+                                          in_runtime=True)
+        elif refers_to_class_or_function(expr.base):
             # Special form -- type application.
             # Translate index to an unanalyzed type.
             types = []  # type: List[Type]
@@ -2375,7 +2400,7 @@ def visit_index_expr(self, expr: IndexExpr) -> None:
                 except TypeTranslationError:
                     self.fail('Type expected within [...]', expr)
                     return
-                typearg = self.anal_type(typearg)
+                typearg = self.anal_type(typearg, aliasing=True)
                 types.append(typearg)
             expr.analyzed = TypeApplication(expr.base, types)
             expr.analyzed.line = expr.line
@@ -3051,6 +3076,8 @@ def visit_decorator(self, dec: Decorator) -> None:
 
     def visit_assignment_stmt(self, s: AssignmentStmt) -> None:
         self.analyze(s.type)
+        if isinstance(s.rvalue, IndexExpr) and isinstance(s.rvalue.analyzed, TypeAliasExpr):
+            self.analyze(s.rvalue.analyzed.type)
         super().visit_assignment_stmt(s)
 
     def visit_cast_expr(self, e: CastExpr) -> None: