fixup! Search sys.path for PEP-561 compliant packages

Author: AWhetter

mypy/main.py

@@ -1033,10 +1033,10 @@ def set_strict_flags() -> None:
     # Set target.
     if special_opts.modules + special_opts.packages:
         options.build_type = BuildType.MODULE
-        egg_dirs, site_packages, sys_path = get_search_dirs(options.python_executable)
+        search_dirs = get_search_dirs(options.python_executable)
         search_paths = SearchPaths((os.getcwd(),),
                                    tuple(mypy_path() + options.mypy_path),
-                                   tuple(egg_dirs + site_packages + sys_path),
+                                   tuple(search_dirs),
                                    ())
         targets = []
         # TODO: use the same cache that the BuildManager will

mypy/modulefinder.py

@@ -608,70 +608,26 @@ def default_lib_path(data_dir: str,
 
 
 @functools.lru_cache(maxsize=None)
-def get_search_dirs(python_executable: Optional[str]) -> Tuple[List[str], List[str], List[str]]:
+def get_search_dirs(python_executable: Optional[str]) -> List[str]:
     """Find package directories for given python.
 
-    This runs a subprocess call, which generates a list of the egg directories, and the site
-    package directories. To avoid repeatedly calling a subprocess (which can be slow!) we
+    This runs a subprocess call, which generates a list of the directories in sys.path.
+    To avoid repeatedly calling a subprocess (which can be slow!) we
     lru_cache the results.
     """
 
     if python_executable is None:
-        return [], [], []
+        return []
     elif python_executable == sys.executable:
         # Use running Python's package dirs
-        site_packages, sys_path = pyinfo.getsearchdirs()
+        sys_path = pyinfo.getsearchdirs()
     else:
         # Use subprocess to get the package directory of given Python
         # executable
-        site_packages, sys_path = ast.literal_eval(
+        sys_path = ast.literal_eval(
             subprocess.check_output([python_executable, pyinfo.__file__, 'getsearchdirs'],
             stderr=subprocess.PIPE).decode())
-    return expand_site_packages(site_packages) + (sys_path,)
-
-
-def expand_site_packages(site_packages: List[str]) -> Tuple[List[str], List[str]]:
-    """Expands .pth imports in site-packages directories"""
-    egg_dirs: List[str] = []
-    for dir in site_packages:
-        if not os.path.isdir(dir):
-            continue
-        pth_filenames = sorted(name for name in os.listdir(dir) if name.endswith(".pth"))
-        for pth_filename in pth_filenames:
-            egg_dirs.extend(_parse_pth_file(dir, pth_filename))
-
-    return egg_dirs, site_packages
-
-
-def _parse_pth_file(dir: str, pth_filename: str) -> Iterator[str]:
-    """
-    Mimics a subset of .pth import hook from Lib/site.py
-    See https://github.com/python/cpython/blob/3.5/Lib/site.py#L146-L185
-    """
-
-    pth_file = os.path.join(dir, pth_filename)
-    try:
-        f = open(pth_file, "r")
-    except OSError:
-        return
-    with f:
-        for line in f.readlines():
-            if line.startswith("#"):
-                # Skip comment lines
-                continue
-            if line.startswith(("import ", "import\t")):
-                # import statements in .pth files are not supported
-                continue
-
-            yield _make_abspath(line.rstrip(), dir)
-
-
-def _make_abspath(path: str, root: str) -> str:
-    """Take a path and make it absolute relative to root if not already absolute."""
-    if os.path.isabs(path):
-        return os.path.normpath(path)
-    else:
-        return os.path.join(root, os.path.normpath(path))
+    return sys_path
 
 
 def add_py2_mypypath_entries(mypypath: List[str]) -> List[str]:
@@ -760,20 +716,20 @@ def compute_search_paths(sources: List[BuildSource],
     if options.python_version[0] == 2:
         mypypath = add_py2_mypypath_entries(mypypath)
 
-    egg_dirs, site_packages, sys_path = get_search_dirs(options.python_executable)
-    for site_dir in site_packages + sys_path:
-        assert site_dir not in lib_path
-        if (site_dir in mypypath or
-                any(p.startswith(site_dir + os.path.sep) for p in mypypath) or
-                os.path.altsep and any(p.startswith(site_dir + os.path.altsep) for p in mypypath)):
-            print("{} is in the MYPYPATH. Please remove it.".format(site_dir), file=sys.stderr)
+    search_dirs = get_search_dirs(options.python_executable)
+    for search_dir in search_dirs:
+        assert search_dir not in lib_path
+        if (search_dir in mypypath or
+                any(p.startswith(search_dir + os.path.sep) for p in mypypath) or
+                os.path.altsep and any(p.startswith(search_dir + os.path.altsep) for p in mypypath)):
+            print("{} is in the MYPYPATH. Please remove it.".format(search_dir), file=sys.stderr)
             print("See https://mypy.readthedocs.io/en/stable/running_mypy.html"
                   "#how-mypy-handles-imports for more info", file=sys.stderr)
             sys.exit(1)
 
     return SearchPaths(python_path=tuple(reversed(python_path)),
                        mypy_path=tuple(mypypath),
-                       package_path=tuple(egg_dirs + site_packages + sys_path),
+                       package_path=tuple(search_dirs),
                        typeshed_path=tuple(lib_path))
 
 

mypy/pyinfo.py

@@ -20,9 +20,7 @@
 
 
 def getsearchdirs():
-    # type: () -> Tuple[List[str], List[str]]
-    site_packages = _getsitepackages()
-
+    # type: () -> List[str]
     # Do not include things from the standard library
     # because those should come from typeshed.
     stdlib_zip = os.path.join(
@@ -33,24 +31,10 @@ def getsearchdirs():
     stdlib = sysconfig.get_path("stdlib")
     stdlib_ext = os.path.join(stdlib, "lib-dynload")
     cwd = os.path.abspath(os.getcwd())
-    excludes = set(site_packages + [cwd, stdlib_zip, stdlib, stdlib_ext])
+    excludes = set([cwd, stdlib_zip, stdlib, stdlib_ext])
 
     abs_sys_path = (os.path.abspath(p) for p in sys.path)
-    return (site_packages, [p for p in abs_sys_path if p not in excludes])
-
-
-def _getsitepackages():
-    # type: () -> List[str]
-    res = []
-    if hasattr(site, 'getsitepackages'):
-        res.extend(site.getsitepackages())
-
-        if hasattr(site, 'getusersitepackages') and site.ENABLE_USER_SITE:
-            res.insert(0, site.getusersitepackages())
-    else:
-        from distutils.sysconfig import get_python_lib
-        res = [get_python_lib()]
-    return res
+    return [p for p in abs_sys_path if p not in excludes]
 
 
 if __name__ == '__main__':

mypy/test/testmodulefinder.py

@@ -1,18 +1,62 @@
 import os
+from typing import Iterator, List, Tuple
 
 from mypy.options import Options
 from mypy.modulefinder import (
     FindModuleCache,
     SearchPaths,
     ModuleNotFoundReason,
-    expand_site_packages
 )
 
 from mypy.test.helpers import Suite, assert_equal
 from mypy.test.config import package_path
 data_path = os.path.relpath(os.path.join(package_path, "modulefinder"))
 
 
+def expand_site_packages(site_packages: List[str]) -> Tuple[List[str], List[str]]:
+    """Mimic the .pth imports in site-packages directories."""
+    egg_dirs: List[str] = []
+    for dir in site_packages:
+        if not os.path.isdir(dir):
+            continue
+        pth_filenames = sorted(name for name in os.listdir(dir) if name.endswith(".pth"))
+        for pth_filename in pth_filenames:
+            egg_dirs.extend(_parse_pth_file(dir, pth_filename))
+
+    return egg_dirs, site_packages
+
+
+def _parse_pth_file(dir: str, pth_filename: str) -> Iterator[str]:
+    """
+    Mimics a subset of .pth import hook from Lib/site.py
+    See https://github.com/python/cpython/blob/3.5/Lib/site.py#L146-L185
+    """
+
+    pth_file = os.path.join(dir, pth_filename)
+    try:
+        f = open(pth_file, "r")
+    except OSError:
+        return
+    with f:
+        for line in f.readlines():
+            if line.startswith("#"):
+                # Skip comment lines
+                continue
+            if line.startswith(("import ", "import\t")):
+                # import statements in .pth files are not supported
+                continue
+
+            yield _make_abspath(line.rstrip(), dir)
+
+
+def _make_abspath(path: str, root: str) -> str:
+    """Take a path and make it absolute relative to root if not already absolute."""
+    if os.path.isabs(path):
+        return os.path.normpath(path)
+    else:
+        return os.path.join(root, os.path.normpath(path))
+
+
 class ModuleFinderSuite(Suite):
 
     def setUp(self) -> None: