Search sys.path for PEP-561 compliant packages (#11143)

Closes #5701

This replaces the old hand crafted search code that was more fragile.

Author: AWhetter

mypy/main.py

@@ -16,7 +16,7 @@
 from mypy import util
 from mypy.modulefinder import (
     BuildSource, FindModuleCache, SearchPaths,
-    get_site_packages_dirs, mypy_path,
+    get_search_dirs, mypy_path,
 )
 from mypy.find_sources import create_source_list, InvalidSourceList
 from mypy.fscache import FileSystemCache
@@ -1043,10 +1043,10 @@ def set_strict_flags() -> None:
     # Set target.
     if special_opts.modules + special_opts.packages:
         options.build_type = BuildType.MODULE
-        egg_dirs, site_packages = get_site_packages_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),
+                                   tuple(search_dirs),
                                    ())
         targets = []
         # TODO: use the same cache that the BuildManager will

mypy/modulefinder.py

@@ -19,7 +19,7 @@
 else:
     import tomli as tomllib
 
-from typing import Dict, Iterator, List, NamedTuple, Optional, Set, Tuple, Union
+from typing import Dict, List, NamedTuple, Optional, Set, Tuple, Union
 from typing_extensions import Final, TypeAlias as _TypeAlias
 
 from mypy.fscache import FileSystemCache
@@ -330,6 +330,9 @@ def _find_module_non_stub_helper(self, components: List[str],
             elif not plausible_match and (self.fscache.isdir(dir_path)
                                           or self.fscache.isfile(dir_path + ".py")):
                 plausible_match = True
+            # If this is not a directory then we can't traverse further into it
+            if not self.fscache.isdir(dir_path):
+                break
         if is_legacy_bundled_package(components[0], self.python_major_ver):
             if (len(components) == 1
                     or (self.find_module(components[0]) is
@@ -724,97 +727,32 @@ def default_lib_path(data_dir: str,
 
 
 @functools.lru_cache(maxsize=None)
-def get_prefixes(python_executable: Optional[str]) -> Tuple[str, str]:
-    """Get the sys.base_prefix and sys.prefix for the given python.
-
-    This runs a subprocess call to get the prefix paths of the given Python executable.
-    To avoid repeatedly calling a subprocess (which can be slow!) we
-    lru_cache the results.
-    """
-    if python_executable is None:
-        return '', ''
-    elif python_executable == sys.executable:
-        # Use running Python's package dirs
-        return pyinfo.getprefixes()
-    else:
-        # Use subprocess to get the package directory of given Python
-        # executable
-        return ast.literal_eval(
-            subprocess.check_output([python_executable, pyinfo.__file__, 'getprefixes'],
-            stderr=subprocess.PIPE).decode())
-
-
-@functools.lru_cache(maxsize=None)
-def get_site_packages_dirs(python_executable: Optional[str]) -> Tuple[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 = pyinfo.getsitepackages()
+        sys_path = pyinfo.getsearchdirs()
     else:
         # Use subprocess to get the package directory of given Python
         # executable
         try:
-            site_packages = ast.literal_eval(
-                subprocess.check_output([python_executable, pyinfo.__file__, 'getsitepackages'],
+            sys_path = ast.literal_eval(
+                subprocess.check_output([python_executable, pyinfo.__file__, 'getsearchdirs'],
                 stderr=subprocess.PIPE).decode())
         except OSError as err:
             reason = os.strerror(err.errno)
             raise CompileError(
                 [f"mypy: Invalid python executable '{python_executable}': {reason}"]
             ) from err
-    return expand_site_packages(site_packages)
-
-
-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)
-    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]:
@@ -903,27 +841,21 @@ def compute_search_paths(sources: List[BuildSource],
     if options.python_version[0] == 2:
         mypypath = add_py2_mypypath_entries(mypypath)
 
-    egg_dirs, site_packages = get_site_packages_dirs(options.python_executable)
-    base_prefix, prefix = get_prefixes(options.python_executable)
-    is_venv = base_prefix != prefix
-    for site_dir in site_packages:
-        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(f"{site_dir} is in the MYPYPATH. Please remove it.", 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(f"{search_dir} is in the MYPYPATH. Please remove it.", 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)
-        elif site_dir in python_path and (is_venv and not site_dir.startswith(prefix)):
-            print("{} is in the PYTHONPATH. Please change directory"
-                  " so it is not.".format(site_dir),
-                  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),
+                       package_path=tuple(search_dirs),
                        typeshed_path=tuple(lib_path))
 
 

mypy/pyinfo.py

@@ -6,41 +6,39 @@
 library found in Python 2. This file is run each mypy run, so it should be kept as fast as
 possible.
 """
-import site
+import os
 import sys
+import sysconfig
 
 if __name__ == '__main__':
     sys.path = sys.path[1:]  # we don't want to pick up mypy.types
 
 MYPY = False
 if MYPY:
-    from typing import List, Tuple
+    from typing import List
 
 
-def getprefixes():
-    # type: () -> Tuple[str, str]
-    return getattr(sys, "base_prefix", sys.prefix), sys.prefix
-
-
-def getsitepackages():
+def getsearchdirs():
     # 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
+    # Do not include things from the standard library
+    # because those should come from typeshed.
+    stdlib_zip = os.path.join(
+        sys.base_exec_prefix,
+        getattr(sys, "platlibdir", "lib"),
+        "python{}{}.zip".format(sys.version_info.major, sys.version_info.minor)
+    )
+    stdlib = sysconfig.get_path("stdlib")
+    stdlib_ext = os.path.join(stdlib, "lib-dynload")
+    cwd = os.path.abspath(os.getcwd())
+    excludes = set([cwd, stdlib_zip, stdlib, stdlib_ext])
+
+    abs_sys_path = (os.path.abspath(p) for p in sys.path)
+    return [p for p in abs_sys_path if p not in excludes]
 
 
 if __name__ == '__main__':
-    if sys.argv[-1] == 'getsitepackages':
-        print(repr(getsitepackages()))
-    elif sys.argv[-1] == 'getprefixes':
-        print(repr(getprefixes()))
+    if sys.argv[-1] == 'getsearchdirs':
+        print(repr(getsearchdirs()))
     else:
         print("ERROR: incorrect argument to pyinfo.py.", file=sys.stderr)
         sys.exit(1)

mypy/test/testcmdline.py

@@ -65,7 +65,10 @@ def test_python_cmdline(testcase: DataDrivenTestCase, step: int) -> None:
     fixed = [python3_path, '-m', 'mypy']
     env = os.environ.copy()
     env.pop('COLUMNS', None)
+    extra_path = os.path.join(os.path.abspath(test_temp_dir), 'pypath')
     env['PYTHONPATH'] = PREFIX
+    if os.path.isdir(extra_path):
+        env['PYTHONPATH'] += os.pathsep + extra_path
     process = subprocess.Popen(fixed + args,
                                stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE,

mypy/test/testmodulefinder.py

@@ -5,7 +5,6 @@
     FindModuleCache,
     SearchPaths,
     ModuleNotFoundReason,
-    expand_site_packages
 )
 
 from mypy.test.helpers import Suite, assert_equal
@@ -149,12 +148,17 @@ def setUp(self) -> None:
             "modulefinder-site-packages",
         ))
 
-        egg_dirs, site_packages = expand_site_packages([self.package_dir])
+        package_paths = (
+            os.path.join(self.package_dir, "baz"),
+            os.path.join(self.package_dir, "..", "not-a-directory"),
+            os.path.join(self.package_dir, "..", "modulefinder-src"),
+            self.package_dir,
+        )
 
         self.search_paths = SearchPaths(
             python_path=(),
             mypy_path=(os.path.join(data_path, "pkg1"),),
-            package_path=tuple(egg_dirs + site_packages),
+            package_path=tuple(package_paths),
             typeshed_path=(),
         )
         options = Options()

test-data/packages/modulefinder-site-packages/baz.pth

@@ -365,6 +365,32 @@ main.py:6: error: Unsupported operand types for + ("int" and "str")
 main.py:7: error: Module has no attribute "y"
 main.py:8: error: Unsupported operand types for + (Module and "int")
 
+[case testConfigFollowImportsSysPath]
+# cmd: mypy main.py
+[file main.py]
+from a import x
+x + 0
+x + ''  # E
+import a
+a.x + 0
+a.x + ''  # E
+a.y  # E
+a + 0  # E
+[file mypy.ini]
+\[mypy]
+follow_imports = normal
+no_silence_site_packages = True
+[file pypath/a/__init__.py]
+x = 0
+x += ''  # Error reported here
+[file pypath/a/py.typed]
+[out]
+pypath/a/__init__.py:2: error: Unsupported operand types for + ("int" and "str")
+main.py:3: error: Unsupported operand types for + ("int" and "str")
+main.py:6: error: Unsupported operand types for + ("int" and "str")
+main.py:7: error: Module has no attribute "y"
+main.py:8: error: Unsupported operand types for + (Module and "int")
+
 [case testConfigFollowImportsSilent]
 # cmd: mypy main.py
 [file main.py]