Add module reference extractor for types

This commit adds a new visitor to extract all module references from
arbitrary types.

This visitor should be considered provisional -- it will most likely be
merged in with the typechecker at a later stage to make it a less
expensive call.

Author: MichaelLeeDBX

mypy/checkexpr.py

@@ -36,6 +36,7 @@
 from mypy.constraints import get_actual_type
 from mypy.checkstrformat import StringFormatterChecker
 from mypy.expandtype import expand_type
+from mypy.util import split_module_names
 
 from mypy import experiments
 
@@ -45,17 +46,6 @@
                       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.
 

mypy/indirection.py

@@ -0,0 +1,101 @@
+from typing import Optional, Set, Iterable
+from abc import abstractmethod
+
+from mypy.visitor import NodeVisitor
+from mypy.types import TypeVisitor
+from mypy.nodes import MODULE_REF
+import mypy.nodes as nodes
+import mypy.types as types
+from mypy.util import split_module_names
+
+
+def extract_module_names(symbol_name: str) -> Iterable[str]:
+    """Returns the module and parent modules of a fully qualified symbol name."""
+    if symbol_name is not None:
+        while '.' in symbol_name:
+            symbol_name = symbol_name.rsplit('.', 1)[0]
+            yield symbol_name
+
+
+class TypeIndirectionVisitor(TypeVisitor[Set[str]]):
+    """Returns all module references within a particular type."""
+
+    def __init__(self) -> None:
+        self.cache = {}  # type: Dict[types.Type, Set[str]]
+
+    def find_modules(self, typs: Iterable[types.Type]) -> Set[str]:
+        return self._visit(*typs)
+
+    def _visit(self, *typs: types.Type) -> Set[str]:
+        output = set()  # type: Set[str]
+        for typ in typs:
+            if typ in self.cache:
+                modules = self.cache[typ]
+            else:
+                modules = typ.accept(self)
+                self.cache[typ] = set(modules)
+            output.update(modules)
+        return output
+
+    def visit_unbound_type(self, t: types.UnboundType) -> Set[str]:
+        return self._visit(*t.args)
+
+    def visit_type_list(self, t: types.TypeList) -> Set[str]:
+        return self._visit(*t.items)
+
+    def visit_error_type(self, t: types.ErrorType) -> Set[str]:
+        return set()
+
+    def visit_any(self, t: types.AnyType) -> Set[str]:
+        return set()
+
+    def visit_void(self, t: types.Void) -> Set[str]:
+        return set()
+
+    def visit_none_type(self, t: types.NoneTyp) -> Set[str]:
+        return set()
+
+    def visit_uninhabited_type(self, t: types.UninhabitedType) -> Set[str]:
+        return set()
+
+    def visit_erased_type(self, t: types.ErasedType) -> Set[str]:
+        return set()
+
+    def visit_deleted_type(self, t: types.DeletedType) -> Set[str]:
+        return set()
+
+    def visit_type_var(self, t: types.TypeVarType) -> Set[str]:
+        return self._visit(*t.values) | self._visit(t.upper_bound)
+
+    def visit_instance(self, t: types.Instance) -> Set[str]:
+        out = self._visit(*t.args)
+        if t.type is not None:
+            out.update(split_module_names(t.type.module_name))
+        return out
+
+    def visit_callable_type(self, t: types.CallableType) -> Set[str]:
+        out = self._visit(*t.arg_types) | self._visit(t.ret_type)
+        if t.definition is not None:
+            out.update(extract_module_names(t.definition.fullname()))
+        return out
+
+    def visit_overloaded(self, t: types.Overloaded) -> Set[str]:
+        return self._visit(*t.items()) | self._visit(t.fallback)
+
+    def visit_tuple_type(self, t: types.TupleType) -> Set[str]:
+        return self._visit(*t.items) | self._visit(t.fallback)
+
+    def visit_star_type(self, t: types.StarType) -> Set[str]:
+        return set()
+
+    def visit_union_type(self, t: types.UnionType) -> Set[str]:
+        return self._visit(*t.items)
+
+    def visit_partial_type(self, t: types.PartialType) -> Set[str]:
+        return set()
+
+    def visit_ellipsis_type(self, t: types.EllipsisType) -> Set[str]:
+        return set()
+
+    def visit_type_type(self, t: types.TypeType) -> Set[str]:
+        return self._visit(t.item)

mypy/util.py

@@ -2,7 +2,7 @@
 
 import re
 import subprocess
-from typing import TypeVar, List, Any, Tuple, Optional
+from typing import TypeVar, List, Any, Tuple, Optional, Iterable
 
 
 T = TypeVar('T')
@@ -12,6 +12,17 @@
 default_python2_interpreter = ['python2', 'python', '/usr/bin/python']
 
 
+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 short_type(obj: object) -> str:
     """Return the last component of the type name of an object.