Modify typechecker to record all module references

This commit modifies the typechecking stage to record all module
references that a particular module makes. This information will be
used in a future commit.

The motivation for this change is that it turns out that the set of
all module references is often larger then the set of explicitly
imported modules, and having this information would allow us to make
incremental mode more aggressive.

Author: MichaelLeeDBX

mypy/checker.py

@@ -86,6 +86,8 @@ class TypeChecker(NodeVisitor[Type]):
     msg = None  # type: MessageBuilder
     # Types of type checked nodes
     type_map = None  # type: Dict[Node, Type]
+    # Types of type checked nodes within this specific module
+    module_type_map = None  # type: Dict[Node, Type]
 
     # Helper for managing conditional types
     binder = None  # type: ConditionalTypeBinder
@@ -117,6 +119,10 @@ class TypeChecker(NodeVisitor[Type]):
     is_typeshed_stub = False
     options = None  # type: Options
 
+    # The set of all dependencies (suppressed or not) that this module accesses, either
+    # directly or indirectly.
+    module_refs = None  # type: Set[str]
+
     def __init__(self, errors: Errors, modules: Dict[str, MypyFile], options: Options) -> None:
         """Construct a type checker.
 
@@ -127,6 +133,7 @@ def __init__(self, errors: Errors, modules: Dict[str, MypyFile], options: Option
         self.options = options
         self.msg = MessageBuilder(errors, modules)
         self.type_map = {}
+        self.module_type_map = {}
         self.binder = ConditionalTypeBinder()
         self.expr_checker = mypy.checkexpr.ExpressionChecker(self, self.msg)
         self.return_types = []
@@ -138,6 +145,7 @@ def __init__(self, errors: Errors, modules: Dict[str, MypyFile], options: Option
         self.deferred_nodes = []
         self.pass_num = 0
         self.current_node_deferred = False
+        self.module_refs = set()
 
     def visit_file(self, file_node: MypyFile, path: str) -> None:
         """Type check a mypy file with the given path."""
@@ -148,6 +156,8 @@ def visit_file(self, file_node: MypyFile, path: str) -> None:
         self.weak_opts = file_node.weak_opts
         self.enter_partial_types()
         self.is_typeshed_stub = self.errors.is_typeshed_file(path)
+        self.module_type_map = {}
+        self.module_refs = set()
 
         for d in file_node.defs:
             self.accept(d)
@@ -2186,6 +2196,8 @@ def check_type_equivalency(self, t1: Type, t2: Type, node: Context,
     def store_type(self, node: Node, typ: Type) -> None:
         """Store the type of a node in the type map."""
         self.type_map[node] = typ
+        if typ is not None:
+            self.module_type_map[node] = typ
 
     def typing_mode_none(self) -> bool:
         if self.is_dynamic_function() and not self.options.check_untyped_defs:

mypy/checkexpr.py

@@ -1,6 +1,6 @@
 """Expression type checker. This file is conceptually part of TypeChecker."""
 
-from typing import cast, Dict, List, Tuple, Callable, Union, Optional
+from typing import cast, Dict, Set, List, Iterable, Tuple, Callable, Union, Optional
 
 from mypy.types import (
     Type, AnyType, CallableType, Overloaded, NoneTyp, Void, TypeVarDef,
@@ -16,7 +16,7 @@
     ListComprehension, GeneratorExpr, SetExpr, MypyFile, Decorator,
     ConditionalExpr, ComparisonExpr, TempNode, SetComprehension,
     DictionaryComprehension, ComplexExpr, EllipsisExpr, StarExpr,
-    TypeAliasExpr, BackquoteExpr, ARG_POS, ARG_NAMED, ARG_STAR2
+    TypeAliasExpr, BackquoteExpr, ARG_POS, ARG_NAMED, ARG_STAR2, MODULE_REF,
 )
 from mypy.nodes import function_type
 from mypy import nodes
@@ -45,6 +45,43 @@
                       None]
 
 
+def split_module_names(mod_name: str) -> Iterable[str]:
+    """Yields the module and all parent module names.
+
+    So, if `mod_name` is 'a.b.c', this function will yield
+    ['a.b.c', 'a.b', and 'a']."""
+    yield mod_name
+    while '.' in mod_name:
+        mod_name = mod_name.rsplit('.', 1)[0]
+        yield mod_name
+
+
+def extract_refexpr_names(expr: RefExpr) -> Set[str]:
+    """Recursively extracts all module references from a reference expression.
+
+    Note that currently, the only two subclasses of RefExpr are NameExpr and
+    MemberExpr."""
+    output = set()  # type: Set[str]
+    while expr.kind == MODULE_REF or expr.fullname is not None:
+        if expr.kind == MODULE_REF:
+            output.add(expr.fullname)
+        elif expr.fullname is not None and '.' in expr.fullname:
+            output.add(expr.fullname.rsplit('.', 1)[0])
+
+        if isinstance(expr, NameExpr):
+            if expr.info is not None:
+                output.update(split_module_names(expr.info.module_name))
+            break
+        elif isinstance(expr, MemberExpr):
+            if isinstance(expr.expr, RefExpr):
+                expr = expr.expr
+            else:
+                break
+        else:
+            raise AssertionError("Unknown RefExpr subclass: {}".format(type(expr)))
+    return output
+
+
 class Finished(Exception):
     """Raised if we can terminate overload argument check early (no match)."""
 
@@ -70,11 +107,16 @@ def __init__(self,
         self.msg = msg
         self.strfrm_checker = StringFormatterChecker(self, self.chk, self.msg)
 
+    def _visit_typeinfo(self, info: nodes.TypeInfo) -> None:
+        if info is not None:
+            self.chk.module_refs.update(split_module_names(info.module_name))
+
     def visit_name_expr(self, e: NameExpr) -> Type:
         """Type check a name expression.
 
         It can be of any kind: local, member or global.
         """
+        self.chk.module_refs.update(extract_refexpr_names(e))
         result = self.analyze_ref_expr(e)
         return self.chk.narrow_type_from_binder(e, result)
 
@@ -858,6 +900,7 @@ def apply_generic_arguments2(self, overload: Overloaded, types: List[Type],
 
     def visit_member_expr(self, e: MemberExpr) -> Type:
         """Visit member expression (of form e.id)."""
+        self.chk.module_refs.update(extract_refexpr_names(e))
         result = self.analyze_ordinary_member_access(e, False)
         return self.chk.narrow_type_from_binder(e, result)