Merge branch 'collection-type-infer'

Partially addresses #1055.

Author: JukkaL

mypy/checker.py

@@ -29,7 +29,7 @@
 from mypy.types import (
     Type, AnyType, CallableType, Void, FunctionLike, Overloaded, TupleType,
     Instance, NoneTyp, UnboundType, ErrorType, TypeTranslator, strip_type,
-    UnionType, TypeVarType,
+    UnionType, TypeVarType, PartialType
 )
 from mypy.sametypes import is_same_type
 from mypy.messages import MessageBuilder
@@ -332,7 +332,8 @@ class TypeChecker(NodeVisitor[Type]):
     breaking_out = False
     # Do weak type checking in this file
     weak_opts = set()        # type: Set[str]
-
+    # Stack of collections of variables with partial types
+    partial_types = None  # type: List[Dict[Var, Context]]
     globals = None  # type: SymbolTable
     locals = None  # type: SymbolTable
     modules = None  # type: Dict[str, MypyFile]
@@ -358,6 +359,7 @@ def __init__(self, errors: Errors, modules: Dict[str, MypyFile],
         self.dynamic_funcs = []
         self.function_stack = []
         self.weak_opts = set()  # type: Set[str]
+        self.partial_types = []
 
     def visit_file(self, file_node: MypyFile, path: str) -> None:
         """Type check a mypy file with the given path."""
@@ -367,10 +369,12 @@ def visit_file(self, file_node: MypyFile, path: str) -> None:
         self.globals = file_node.names
         self.locals = None
         self.weak_opts = file_node.weak_opts
+        self.enter_partial_types()
 
         for d in file_node.defs:
             self.accept(d)
 
+        self.leave_partial_types()
         self.errors.set_ignored_lines(set())
 
     def accept(self, node: Node, type_context: Type = None) -> Type:
@@ -461,6 +465,8 @@ def check_func_item(self, defn: FuncItem,
         if fdef:
             self.errors.push_function(fdef.name())
 
+        self.enter_partial_types()
+
         typ = self.function_type(defn)
         if type_override:
             typ = type_override
@@ -469,6 +475,8 @@ def check_func_item(self, defn: FuncItem,
         else:
             raise RuntimeError('Not supported')
 
+        self.leave_partial_types()
+
         if fdef:
             self.errors.pop_function()
 
@@ -864,12 +872,14 @@ def visit_class_def(self, defn: ClassDef) -> Type:
         """Type check a class definition."""
         typ = defn.info
         self.errors.push_type(defn.name)
+        self.enter_partial_types()
         old_binder = self.binder
         self.binder = ConditionalTypeBinder()
         self.binder.push_frame()
         self.accept(defn.defs)
         self.binder = old_binder
         self.check_multiple_inheritance(typ)
+        self.leave_partial_types()
         self.errors.pop_type()
 
     def check_multiple_inheritance(self, typ: TypeInfo) -> None:
@@ -1237,11 +1247,14 @@ def infer_variable_type(self, name: Var, lvalue: Node,
         elif isinstance(init_type, Void):
             self.check_not_void(init_type, context)
             self.set_inference_error_fallback_type(name, lvalue, init_type, context)
-        elif not self.is_valid_inferred_type(init_type):
-            # We cannot use the type of the initialization expression for type
-            # inference (it's not specific enough).
-            self.fail(messages.NEED_ANNOTATION_FOR_VAR, context)
-            self.set_inference_error_fallback_type(name, lvalue, init_type, context)
+        elif not is_valid_inferred_type(init_type):
+            # We cannot use the type of the initialization expression for full type
+            # inference (it's not specific enough), but we might be able to give
+            # partial type which will be made more specific later. A partial type
+            # gets generated in assignment like 'x = []' where item type is not known.
+            if not self.infer_partial_type(name, lvalue, init_type):
+                self.fail(messages.NEED_ANNOTATION_FOR_VAR, context)
+                self.set_inference_error_fallback_type(name, lvalue, init_type, context)
         else:
             # Infer type of the target.
 
@@ -1250,6 +1263,21 @@ def infer_variable_type(self, name: Var, lvalue: Node,
 
             self.set_inferred_type(name, lvalue, init_type)
 
+    def infer_partial_type(self, name: Var, lvalue: Node, init_type: Type) -> bool:
+        if not isinstance(init_type, Instance):
+            return False
+        fullname = init_type.type.fullname()
+        if ((fullname == 'builtins.list' or fullname == 'builtins.set' or
+             fullname == 'builtins.dict')
+                and isinstance(init_type.args[0], NoneTyp)
+                and (fullname != 'builtins.dict' or isinstance(init_type.args[1], NoneTyp))
+                and isinstance(lvalue, NameExpr)):
+            partial_type = PartialType(init_type.type, name)
+            self.set_inferred_type(name, lvalue, partial_type)
+            self.partial_types[-1][name] = lvalue
+            return True
+        return False
+
     def set_inferred_type(self, var: Var, lvalue: Node, type: Type) -> None:
         """Store inferred variable type.
 
@@ -1275,23 +1303,6 @@ def set_inference_error_fallback_type(self, var: Var, lvalue: Node, type: Type,
         if context.get_line() in self.errors.ignored_lines:
             self.set_inferred_type(var, lvalue, AnyType())
 
-    def is_valid_inferred_type(self, typ: Type) -> bool:
-        """Is an inferred type invalid?
-
-        Examples include the None type or a type with a None component.
-        """
-        if is_same_type(typ, NoneTyp()):
-            return False
-        elif isinstance(typ, Instance):
-            for arg in typ.args:
-                if not self.is_valid_inferred_type(arg):
-                    return False
-        elif isinstance(typ, TupleType):
-            for item in typ.items:
-                if not self.is_valid_inferred_type(item):
-                    return False
-        return True
-
     def narrow_type_from_binder(self, expr: Node, known_type: Type) -> Type:
         if expr.literal >= LITERAL_TYPE:
             restriction = self.binder.get(expr)
@@ -1323,6 +1334,7 @@ def check_indexed_assignment(self, lvalue: IndexExpr,
 
         The lvalue argument is the base[index] expression.
         """
+        self.try_infer_partial_type_from_indexed_assignment(lvalue, rvalue)
         basetype = self.accept(lvalue.base)
         method_type = self.expr_checker.analyze_external_member_access(
             '__setitem__', basetype, context)
@@ -1331,6 +1343,26 @@ def check_indexed_assignment(self, lvalue: IndexExpr,
                                      [nodes.ARG_POS, nodes.ARG_POS],
                                      context)
 
+    def try_infer_partial_type_from_indexed_assignment(
+            self, lvalue: IndexExpr, rvalue: Node) -> None:
+        # TODO: Should we share some of this with try_infer_partial_type?
+        partial_types = self.partial_types[-1]
+        if not partial_types:
+            # Fast path leave -- no partial types in the current scope.
+            return
+        if isinstance(lvalue.base, RefExpr):
+            var = lvalue.base.node
+            if var in partial_types:
+                var = cast(Var, var)
+                typename = cast(Instance, var.type).type.fullname()
+                if typename == 'builtins.dict':
+                    # TODO: Don't infer things twice.
+                    key_type = self.accept(lvalue.index)
+                    value_type = self.accept(rvalue)
+                    if is_valid_inferred_type(key_type) and is_valid_inferred_type(value_type):
+                        var.type = self.named_generic_type('builtins.dict', [key_type, value_type])
+                        del partial_types[var]
+
     def visit_expression_stmt(self, s: ExpressionStmt) -> Type:
         self.accept(s.expr)
 
@@ -2032,6 +2064,21 @@ def enter(self) -> None:
     def leave(self) -> None:
         self.locals = None
 
+    def enter_partial_types(self) -> None:
+        """Push a new scope for collecting partial types."""
+        self.partial_types.append({})
+
+    def leave_partial_types(self) -> None:
+        """Pop partial type scope.
+
+        Also report errors for variables which still have partial
+        types, i.e. we couldn't infer a complete type.
+        """
+        partial_types = self.partial_types.pop()
+        for var, context in partial_types.items():
+            self.msg.fail(messages.NEED_ANNOTATION_FOR_VAR, context)
+            var.type = AnyType()
+
     def is_within_function(self) -> bool:
         """Are we currently type checking within a function?
 
@@ -2289,3 +2336,21 @@ def infer_operator_assignment_method(type: Type, operator: str) -> str:
             if type.type.has_readable_member(inplace):
                 method = inplace
     return method
+
+
+def is_valid_inferred_type(typ: Type) -> bool:
+    """Is an inferred type valid?
+
+    Examples of invalid types include the None type or a type with a None component.
+    """
+    if is_same_type(typ, NoneTyp()):
+        return False
+    elif isinstance(typ, Instance):
+        for arg in typ.args:
+            if not is_valid_inferred_type(arg):
+                return False
+    elif isinstance(typ, TupleType):
+        for item in typ.items:
+            if not is_valid_inferred_type(item):
+                return False
+    return True

mypy/checkexpr.py

@@ -4,7 +4,8 @@
 
 from mypy.types import (
     Type, AnyType, CallableType, Overloaded, NoneTyp, Void, TypeVarDef,
-    TupleType, Instance, TypeVarType, TypeTranslator, ErasedType, FunctionLike, UnionType
+    TupleType, Instance, TypeVarType, TypeTranslator, ErasedType, FunctionLike, UnionType,
+    PartialType
 )
 from mypy.nodes import (
     NameExpr, RefExpr, Var, FuncDef, OverloadedFuncDef, TypeInfo, CallExpr,
@@ -14,7 +15,7 @@
     ListComprehension, GeneratorExpr, SetExpr, MypyFile, Decorator,
     ConditionalExpr, ComparisonExpr, TempNode, SetComprehension,
     DictionaryComprehension, ComplexExpr, EllipsisExpr, LITERAL_TYPE,
-    TypeAliasExpr, YieldExpr, BackquoteExpr
+    TypeAliasExpr, YieldExpr, BackquoteExpr, ARG_POS
 )
 from mypy.errors import Errors
 from mypy.nodes import function_type
@@ -72,6 +73,11 @@ def analyze_ref_expr(self, e: RefExpr) -> Type:
         if isinstance(node, Var):
             # Variable reference.
             result = self.analyze_var_ref(node, e)
+            if isinstance(result, PartialType):
+                partial_types = self.chk.partial_types[-1]
+                context = partial_types[node]
+                self.msg.fail(messages.NEED_ANNOTATION_FOR_VAR, context)
+                result = AnyType()
         elif isinstance(node, FuncDef):
             # Reference to a global function.
             result = function_type(node, self.named_type('builtins.function'))
@@ -110,13 +116,35 @@ def visit_call_expr(self, e: CallExpr) -> Type:
         if e.analyzed:
             # It's really a special form that only looks like a call.
             return self.accept(e.analyzed)
+        self.try_infer_partial_type(e)
         self.accept(e.callee)
         # Access callee type directly, since accept may return the Any type
         # even if the type is known (in a dynamically typed function). This
         # way we get a more precise callee in dynamically typed functions.
         callee_type = self.chk.type_map[e.callee]
         return self.check_call_expr_with_callee_type(callee_type, e)
 
+    def try_infer_partial_type(self, e: CallExpr) -> None:
+        partial_types = self.chk.partial_types[-1]
+        if not partial_types:
+            # Fast path leave -- no partial types in the current scope.
+            return
+        if isinstance(e.callee, MemberExpr) and isinstance(e.callee.expr, RefExpr):
+            var = e.callee.expr.node
+            if var in partial_types:
+                var = cast(Var, var)
+                typename = cast(Instance, var.type).type.fullname()
+                methodname = e.callee.name
+                if (((typename == 'builtins.list' and methodname == 'append') or
+                     (typename == 'builtins.set' and methodname == 'add'))
+                        and e.arg_kinds == [ARG_POS]):
+                    # We can infer a full type for a partial List type.
+                    # TODO: Don't infer argument expression twice.
+                    item_type = self.accept(e.args[0])
+                    if mypy.checker.is_valid_inferred_type(item_type):
+                        var.type = self.chk.named_generic_type(typename, [item_type])
+                        del partial_types[var]
+
     def check_call_expr_with_callee_type(self, callee_type: Type,
                                          e: CallExpr) -> Type:
         """Type check call expression.

mypy/checkmember.py

@@ -4,7 +4,7 @@
 
 from mypy.types import (
     Type, Instance, AnyType, TupleType, CallableType, FunctionLike, TypeVarDef,
-    Overloaded, TypeVarType, TypeTranslator, UnionType
+    Overloaded, TypeVarType, TypeTranslator, UnionType, PartialType
 )
 from mypy.nodes import TypeInfo, FuncBase, Var, FuncDef, SymbolNode, Context
 from mypy.nodes import ARG_POS, function_type, Decorator, OverloadedFuncDef

mypy/constraints.py

@@ -4,7 +4,7 @@
 
 from mypy.types import (
     CallableType, Type, TypeVisitor, UnboundType, AnyType, Void, NoneTyp, TypeVarType,
-    Instance, TupleType, UnionType, Overloaded, ErasedType, is_named_instance
+    Instance, TupleType, UnionType, Overloaded, ErasedType, PartialType, is_named_instance
 )
 from mypy.expandtype import expand_caller_var_args
 from mypy.maptype import map_instance_to_supertype
@@ -151,6 +151,12 @@ def visit_none_type(self, template: NoneTyp) -> List[Constraint]:
     def visit_erased_type(self, template: ErasedType) -> List[Constraint]:
         return []
 
+    # Errors
+
+    def visit_partial_type(self, template: PartialType) -> List[Constraint]:
+        # We can't do anything useful with a partial type here.
+        assert False, "Internal error"
+
     # Non-trivial leaf type
 
     def visit_type_var(self, template: TypeVarType) -> List[Constraint]:

mypy/erasetype.py

@@ -3,7 +3,7 @@
 from mypy.types import (
     Type, TypeVisitor, UnboundType, ErrorType, AnyType, Void, NoneTyp,
     Instance, TypeVarType, CallableType, TupleType, UnionType, Overloaded, ErasedType,
-    TypeTranslator, TypeList
+    TypeTranslator, TypeList, PartialType
 )
 
 
@@ -46,6 +46,10 @@ def visit_erased_type(self, t: ErasedType) -> Type:
         # Should not get here.
         raise RuntimeError()
 
+    def visit_partial_type(self, t: PartialType) -> Type:
+        # Should not get here.
+        raise RuntimeError()
+
     def visit_instance(self, t: Instance) -> Type:
         return Instance(t.type, [AnyType()] * len(t.args), t.line)
 

mypy/expandtype.py

@@ -2,7 +2,8 @@
 
 from mypy.types import (
     Type, Instance, CallableType, TypeVisitor, UnboundType, ErrorType, AnyType,
-    Void, NoneTyp, TypeVarType, Overloaded, TupleType, UnionType, ErasedType, TypeList
+    Void, NoneTyp, TypeVarType, Overloaded, TupleType, UnionType, ErasedType, TypeList,
+    PartialType
 )
 
 
@@ -92,6 +93,9 @@ def visit_tuple_type(self, t: TupleType) -> Type:
     def visit_union_type(self, t: UnionType) -> Type:
         return UnionType(self.expand_types(t.items), t.line)
 
+    def visit_partial_type(self, t: PartialType) -> Type:
+        return t
+
     def expand_types(self, types: List[Type]) -> List[Type]:
         a = []  # type: List[Type]
         for t in types:

mypy/join.py

@@ -5,7 +5,7 @@
 from mypy.types import (
     Type, AnyType, NoneTyp, Void, TypeVisitor, Instance, UnboundType,
     ErrorType, TypeVarType, CallableType, TupleType, ErasedType, TypeList,
-    UnionType, FunctionLike, Overloaded
+    UnionType, FunctionLike, Overloaded, PartialType
 )
 from mypy.maptype import map_instance_to_supertype
 from mypy.subtypes import is_subtype, is_equivalent, is_subtype_ignoring_tvars
@@ -196,6 +196,11 @@ def visit_tuple_type(self, t: TupleType) -> Type:
         else:
             return self.default(self.s)
 
+    def visit_partial_type(self, t: PartialType) -> Type:
+        # We only have partial information so we can't decide the join result. We should
+        # never get here.
+        assert False, "Internal error"
+
     def join(self, s: Type, t: Type) -> Type:
         return join_types(s, t)