Refactor and reorder search path (#5256)

emmatyping · ilevkivskyi · commit 59947f871d9b · 2018-06-25T22:45:42.000+01:00
As promised in #5227, here is an implementation for refactoring and making the search path compliant with PEP 561. The order is specified in https://www.python.org/dev/peps/pep-0561/#type-checker-module-resolution-order.
diff --git a/docs/source/command_line.rst b/docs/source/command_line.rst
@@ -152,6 +152,8 @@ This is computed from the following items:
   (a colon-separated list of directories).
 - The directories containing the sources given on the command line
   (see below).
+- The installed packages marked as safe for type checking (see
+  :ref:`PEP 561 support <installed-packages>`)
 - The relevant directories of the
   `typeshed <https://github.com/python/typeshed>`_ repo.
 
@@ -161,7 +163,7 @@ contain an ``__init__.py`` or ``__init__.pyi`` file.
 
 Second, mypy searches for stub files in addition to regular Python files
 and packages.
-The rules for searching a module ``foo`` are as follows:
+The rules for searching for a module ``foo`` are as follows:
 
 - The search looks in each of the directories in the search path
   (see above) until a match is found.
diff --git a/mypy/build.py b/mypy/build.py
@@ -188,10 +188,46 @@ def default_flush_errors(new_messages: List[str], is_serious: bool) -> None:
         raise
 
 
-def compute_lib_path(sources: List[BuildSource],
+# python_path is usercode, mypy_path is set via config or envionment variable,
+# package_path is calculated by _get_site_packages_dirs, and typeshed_path points
+# to typeshed. Each is a tuple of paths to be searched in find_module()
+SearchPaths = NamedTuple('SearchPaths',
+             (('python_path', Tuple[str, ...]),
+              ('mypy_path', Tuple[str, ...]),
+              ('package_path', Tuple[str, ...]),
+              ('typeshed_path', Tuple[str, ...])))
+
+
+@functools.lru_cache(maxsize=None)
+def _get_site_packages_dirs(python_executable: Optional[str]) -> List[str]:
+    """Find package directories for given python.
+
+    This runs a subprocess call, which generates a list of the site package directories.
+    To avoid repeatedly calling a subprocess (which can be slow!) we lru_cache the results."""
+    if python_executable is None:
+        return []
+    if python_executable == sys.executable:
+        # Use running Python's package dirs
+        return sitepkgs.getsitepackages()
+    else:
+        # Use subprocess to get the package directory of given Python
+        # executable
+        return ast.literal_eval(subprocess.check_output([python_executable, sitepkgs.__file__],
+                                stderr=subprocess.PIPE).decode())
+
+
+def compute_search_paths(sources: List[BuildSource],
                      options: Options,
                      data_dir: str,
-                     alt_lib_path: Optional[str] = None) -> List[str]:
+                     alt_lib_path: Optional[str] = None) -> SearchPaths:
+    """Compute the search paths as specified in PEP 561.
+
+    There are the following 4 members created:
+    - User code (from `sources`)
+    - MYPYPATH (set either via config or environment variable)
+    - installed package directories (which will later be split into stub-only and inline)
+    - typeshed
+     """
     # Determine the default module search path.
     lib_path = collections.deque(
         default_lib_path(data_dir,
@@ -206,15 +242,14 @@ def compute_lib_path(sources: List[BuildSource],
         lib_path.appendleft(os.path.join(root_dir, 'test-data', 'unit', 'lib-stub'))
     # alt_lib_path is used by some tests to bypass the normal lib_path mechanics.
     # If we don't have one, grab directories of source files.
-    lib_path_set = set(lib_path)
+    python_path = []  # type: List[str]
     if not alt_lib_path:
         for source in sources:
             # Include directory of the program file in the module search path.
             if source.base_dir:
                 dir = source.base_dir
-                if dir not in lib_path_set:
-                    lib_path.appendleft(dir)
-                    lib_path_set.add(dir)
+                if dir not in python_path:
+                    python_path.append(dir)
 
         # Do this even if running as a file, for sanity (mainly because with
         # multiple builds, there could be a mix of files/modules, so its easier
@@ -227,20 +262,23 @@ def compute_lib_path(sources: List[BuildSource],
         else:
             dir = os.getcwd()
         if dir not in lib_path:
-            lib_path.appendleft(dir)
+            python_path.insert(0, dir)
 
-    # Prepend a config-defined mypy path.
-    lib_path.extendleft(options.mypy_path)
+    # Start with a MYPYPATH environment variable at the front of the mypy_path, if defined.
+    mypypath = mypy_path()
 
-    # Add MYPYPATH environment variable to front of library path, if defined.
-    lib_path.extendleft(mypy_path())
+    # Add a config-defined mypy path.
+    mypypath.extend(options.mypy_path)
 
     # If provided, insert the caller-supplied extra module path to the
     # beginning (highest priority) of the search path.
     if alt_lib_path:
-        lib_path.appendleft(alt_lib_path)
+        mypypath.insert(0, alt_lib_path)
 
-    return list(lib_path)
+    return SearchPaths(tuple(reversed(python_path)),
+                       tuple(mypypath),
+                       tuple(_get_site_packages_dirs(options.python_executable)),
+                       tuple(lib_path))
 
 
 def _build(sources: List[BuildSource],
@@ -256,7 +294,7 @@ def _build(sources: List[BuildSource],
     data_dir = default_data_dir(bin_dir)
     fscache = fscache or FileSystemCache()
 
-    lib_path = compute_lib_path(sources, options, data_dir, alt_lib_path)
+    search_paths = compute_search_paths(sources, options, data_dir, alt_lib_path)
 
     reports = Reports(data_dir, options.report_dirs)
     source_set = BuildSourceSet(sources)
@@ -266,7 +304,7 @@ def _build(sources: List[BuildSource],
     # Construct a build manager object to hold state during the build.
     #
     # Ignore current directory prefix in error messages.
-    manager = BuildManager(data_dir, lib_path,
+    manager = BuildManager(data_dir, search_paths,
                            ignore_prefix=os.getcwd(),
                            source_set=source_set,
                            reports=reports,
@@ -613,7 +651,7 @@ class BuildManager:
     """
 
     def __init__(self, data_dir: str,
-                 lib_path: List[str],
+                 search_paths: SearchPaths,
                  ignore_prefix: str,
                  source_set: BuildSourceSet,
                  reports: Reports,
@@ -628,7 +666,7 @@ def __init__(self, data_dir: str,
         self.data_dir = data_dir
         self.errors = errors
         self.errors.set_ignore_prefix(ignore_prefix)
-        self.lib_path = tuple(lib_path)
+        self.search_paths = search_paths
         self.source_set = source_set
         self.reports = reports
         self.options = options
@@ -637,7 +675,7 @@ def __init__(self, data_dir: str,
         self.missing_modules = set()  # type: Set[str]
         self.plugin = plugin
         self.semantic_analyzer = SemanticAnalyzerPass2(self.modules, self.missing_modules,
-                                                  lib_path, self.errors, self.plugin)
+                                                  self.errors, self.plugin)
         self.semantic_analyzer_pass3 = SemanticAnalyzerPass3(self.modules, self.errors,
                                                              self.semantic_analyzer)
         self.all_types = {}  # type: Dict[Expression, Type]  # Used by tests only
@@ -780,7 +818,7 @@ def correct_rel_imp(imp: Union[ImportFrom, ImportAll]) -> str:
 
     def is_module(self, id: str) -> bool:
         """Is there a file in the file system corresponding to module id?"""
-        return self.find_module_cache.find_module(id, self.lib_path,
+        return self.find_module_cache.find_module(id, self.search_paths,
                                                   self.options.python_executable) is not None
 
     def parse_file(self, id: str, path: str, source: str, ignore_errors: bool) -> MypyFile:
@@ -844,26 +882,8 @@ def stats_summary(self) -> Mapping[str, object]:
         return self.stats
 
 
-@functools.lru_cache(maxsize=None)
-def _get_site_packages_dirs(python_executable: Optional[str]) -> List[str]:
-    """Find package directories for given python.
-
-    This runs a subprocess call, which generates a list of the site package directories.
-    To avoid repeatedly calling a subprocess (which can be slow!) we lru_cache the results."""
-    if python_executable is None:
-        return []
-    if python_executable == sys.executable:
-        # Use running Python's package dirs
-        return sitepkgs.getsitepackages()
-    else:
-        # Use subprocess to get the package directory of given Python
-        # executable
-        return ast.literal_eval(subprocess.check_output([python_executable, sitepkgs.__file__],
-                                stderr=subprocess.PIPE).decode())
-
-
-# Search paths are a two-tuple of path and whether to verify the module
-SearchPaths = List[Tuple[str, bool]]
+# Package dirs are a two-tuple of path to search and whether to verify the module
+PackageDirs = List[Tuple[str, bool]]
 
 
 class FindModuleCache:
@@ -879,26 +899,27 @@ class FindModuleCache:
 
     def __init__(self, fscache: Optional[FileSystemCache] = None) -> None:
         self.fscache = fscache or FileSystemCache()
-        # Cache find_lib_path_dirs: (dir_chain, lib_path) -> list of (package_path, should_verify)
-        self.dirs = {}  # type: Dict[Tuple[str, Tuple[str, ...]], SearchPaths]
-        # Cache find_module: (id, lib_path, python_version) -> result.
-        self.results = {}  # type: Dict[Tuple[str, Tuple[str, ...], Optional[str]], Optional[str]]
+        # Cache find_lib_path_dirs: (dir_chain, search_paths) -> list(package_dirs, should_verify)
+        self.dirs = {}  # type: Dict[Tuple[str, Tuple[str, ...]], PackageDirs]
+        # Cache find_module: (id, search_paths, python_version) -> result.
+        self.results = {}  # type: Dict[Tuple[str, SearchPaths, Optional[str]], Optional[str]]
 
     def clear(self) -> None:
         self.results.clear()
         self.dirs.clear()
 
-    def find_lib_path_dirs(self, dir_chain: str, lib_path: Tuple[str, ...]) -> SearchPaths:
+    def find_lib_path_dirs(self, dir_chain: str, lib_path: Tuple[str, ...]) -> PackageDirs:
         # Cache some repeated work within distinct find_module calls: finding which
         # elements of lib_path have even the subdirectory they'd need for the module
-        # to exist.  This is shared among different module ids when they differ only
+        # to exist. This is shared among different module ids when they differ only
         # in the last component.
+        # This is run for the python_path, mypy_path, and typeshed_path search paths
         key = (dir_chain, lib_path)
         if key not in self.dirs:
             self.dirs[key] = self._find_lib_path_dirs(dir_chain, lib_path)
         return self.dirs[key]
 
-    def _find_lib_path_dirs(self, dir_chain: str, lib_path: Tuple[str, ...]) -> SearchPaths:
+    def _find_lib_path_dirs(self, dir_chain: str, lib_path: Tuple[str, ...]) -> PackageDirs:
         dirs = []
         for pathitem in lib_path:
             # e.g., '/usr/lib/python3.4/foo/bar'
@@ -907,15 +928,15 @@ def _find_lib_path_dirs(self, dir_chain: str, lib_path: Tuple[str, ...]) -> Sear
                 dirs.append((dir, True))
         return dirs
 
-    def find_module(self, id: str, lib_path: Tuple[str, ...],
+    def find_module(self, id: str, search_paths: SearchPaths,
                     python_executable: Optional[str]) -> Optional[str]:
         """Return the path of the module source file, or None if not found."""
-        key = (id, lib_path, python_executable)
+        key = (id, search_paths, python_executable)
         if key not in self.results:
-            self.results[key] = self._find_module(id, lib_path, python_executable)
+            self.results[key] = self._find_module(id, search_paths, python_executable)
         return self.results[key]
 
-    def _find_module(self, id: str, lib_path: Tuple[str, ...],
+    def _find_module(self, id: str, search_paths: SearchPaths,
                      python_executable: Optional[str]) -> Optional[str]:
         fscache = self.fscache
 
@@ -932,7 +953,7 @@ def _find_module(self, id: str, lib_path: Tuple[str, ...],
         third_party_inline_dirs = []
         third_party_stubs_dirs = []
         # Third-party stub/typed packages
-        for pkg_dir in _get_site_packages_dirs(python_executable):
+        for pkg_dir in search_paths.package_path:
             stub_name = components[0] + '-stubs'
             typed_file = os.path.join(pkg_dir, components[0], 'py.typed')
             stub_dir = os.path.join(pkg_dir, stub_name)
@@ -957,8 +978,10 @@ def _find_module(self, id: str, lib_path: Tuple[str, ...],
             elif fscache.isfile(typed_file):
                 path = os.path.join(pkg_dir, dir_chain)
                 third_party_inline_dirs.append((path, True))
-        candidate_base_dirs = self.find_lib_path_dirs(dir_chain, lib_path) + \
-            third_party_stubs_dirs + third_party_inline_dirs
+        python_mypy_path = search_paths.python_path + search_paths.mypy_path
+        candidate_base_dirs = self.find_lib_path_dirs(dir_chain, python_mypy_path) + \
+            third_party_stubs_dirs + third_party_inline_dirs + \
+            self.find_lib_path_dirs(dir_chain, search_paths.typeshed_path)
 
         # If we're looking for a module like 'foo.bar.baz', then candidate_base_dirs now
         # contains just the subdirectories 'foo/bar' that actually exist under the
@@ -989,9 +1012,9 @@ def _find_module(self, id: str, lib_path: Tuple[str, ...],
                     return path
         return None
 
-    def find_modules_recursive(self, module: str, lib_path: Tuple[str, ...],
+    def find_modules_recursive(self, module: str, search_paths: SearchPaths,
                                python_executable: Optional[str]) -> List[BuildSource]:
-        module_path = self.find_module(module, lib_path, python_executable)
+        module_path = self.find_module(module, search_paths, python_executable)
         if not module_path:
             return []
         result = [BuildSource(module_path, module, None)]
@@ -1011,14 +1034,14 @@ def find_modules_recursive(self, module: str, lib_path: Tuple[str, ...],
                         (os.path.isfile(os.path.join(abs_path, '__init__.py')) or
                         os.path.isfile(os.path.join(abs_path, '__init__.pyi'))):
                     hits.add(item)
-                    result += self.find_modules_recursive(module + '.' + item, lib_path,
+                    result += self.find_modules_recursive(module + '.' + item, search_paths,
                                                           python_executable)
                 elif item != '__init__.py' and item != '__init__.pyi' and \
                         item.endswith(('.py', '.pyi')):
                     mod = item.split('.')[0]
                     if mod not in hits:
                         hits.add(mod)
-                        result += self.find_modules_recursive(module + '.' + mod, lib_path,
+                        result += self.find_modules_recursive(module + '.' + mod, search_paths,
                                                               python_executable)
         return result
 
@@ -2304,7 +2327,7 @@ def find_module_and_diagnose(manager: BuildManager,
         # difference and just assume 'builtins' everywhere,
         # which simplifies code.
         file_id = '__builtin__'
-    path = manager.find_module_cache.find_module(file_id, manager.lib_path,
+    path = manager.find_module_cache.find_module(file_id, manager.search_paths,
                                                  manager.options.python_executable)
     if path:
         # For non-stubs, look at options.follow_imports:
diff --git a/mypy/dmypy_server.py b/mypy/dmypy_server.py
@@ -326,8 +326,8 @@ def fine_grained_increment(self, sources: List[mypy.build.BuildSource]) -> Dict[
         t0 = time.time()
         self.update_sources(sources)
         changed, removed = self.find_changed(sources)
-        manager.lib_path = tuple(mypy.build.compute_lib_path(
-            sources, manager.options, manager.data_dir))
+        manager.search_paths = mypy.build.compute_search_paths(
+            sources, manager.options, manager.data_dir)
         t1 = time.time()
         messages = self.fine_grained_manager.update(changed, removed)
         t2 = time.time()
diff --git a/mypy/main.py b/mypy/main.py
@@ -15,7 +15,7 @@
 from mypy import defaults
 from mypy import experiments
 from mypy import util
-from mypy.build import BuildSource, BuildResult
+from mypy.build import BuildSource, BuildResult, SearchPaths
 from mypy.find_sources import create_source_list, InvalidSourceList
 from mypy.fscache import FileSystemCache
 from mypy.errors import CompileError
@@ -776,14 +776,14 @@ def add_invertible_flag(flag: str,
     # Set target.
     if special_opts.modules + special_opts.packages:
         options.build_type = BuildType.MODULE
-        lib_path = [os.getcwd()] + build.mypy_path()
+        search_paths = SearchPaths((os.getcwd(),), tuple(build.mypy_path()), (), ())
         targets = []
         # TODO: use the same cache that the BuildManager will
         cache = build.FindModuleCache(fscache)
         for p in special_opts.packages:
             if os.sep in p or os.altsep and os.altsep in p:
                 fail("Package name '{}' cannot have a slash in it.".format(p))
-            p_targets = cache.find_modules_recursive(p, tuple(lib_path), options.python_executable)
+            p_targets = cache.find_modules_recursive(p, search_paths, options.python_executable)
             if not p_targets:
                 fail("Can't find package '{}'".format(p))
             targets.extend(p_targets)
diff --git a/mypy/semanal.py b/mypy/semanal.py
@@ -182,8 +182,6 @@ class SemanticAnalyzerPass2(NodeVisitor[None],
     This is the second phase of semantic analysis.
     """
 
-    # Library search paths
-    lib_path = None  # type: List[str]
     # Module name space
     modules = None  # type: Dict[str, MypyFile]
     # Global name space for current module
@@ -227,7 +225,7 @@ class SemanticAnalyzerPass2(NodeVisitor[None],
     def __init__(self,
                  modules: Dict[str, MypyFile],
                  missing_modules: Set[str],
-                 lib_path: List[str], errors: Errors,
+                 errors: Errors,
                  plugin: Plugin) -> None:
         """Construct semantic analyzer.
 
@@ -242,7 +240,6 @@ def __init__(self,
         self.function_stack = []
         self.block_depth = [0]
         self.loop_depth = 0
-        self.lib_path = lib_path
         self.errors = errors
         self.modules = modules
         self.msg = MessageBuilder(errors, modules)
diff --git a/mypy/stubgen.py b/mypy/stubgen.py
@@ -160,7 +160,8 @@ def find_module_path_and_all(module: str, pyversion: Tuple[int, int],
             module_all = getattr(mod, '__all__', None)
     else:
         # Find module by going through search path.
-        module_path = mypy.build.FindModuleCache().find_module(module, ('.',) + tuple(search_path),
+        search_paths = mypy.build.SearchPaths(('.',) + tuple(search_path), (), (), ())
+        module_path = mypy.build.FindModuleCache().find_module(module, search_paths,
                                                                interpreter)
         if not module_path:
             raise SystemExit(
diff --git a/mypy/test/testcheck.py b/mypy/test/testcheck.py
diff --git a/mypy/test/testgraph.py b/mypy/test/testgraph.py
