Search sys.path for PEP-561 compliant packages

Closes python#5701

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
@@ -1019,10 +1019,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)
+        egg_dirs, site_packages, sys_path = get_search_dirs(options.python_executable)
         search_paths = SearchPaths((os.getcwd(),),
                                    tuple(mypy_path() + options.mypy_path),
-                                   tuple(egg_dirs + site_packages),
+                                   tuple(egg_dirs + site_packages + sys_path),
                                    ())
         targets = []
         # TODO: use the same cache that the BuildManager will

mypy/modulefinder.py

@@ -25,7 +25,7 @@
     'SearchPaths',
     [('python_path', Tuple[str, ...]),  # where user code is found
      ('mypy_path', Tuple[str, ...]),  # from $MYPYPATH or config variable
-     ('package_path', Tuple[str, ...]),  # from get_site_packages_dirs()
+     ('package_path', Tuple[str, ...]),  # from get_search_dirs()
      ('typeshed_path', Tuple[str, ...]),  # paths in typeshed
      ])
 
@@ -585,28 +585,7 @@ 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]) -> Tuple[List[str], List[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
@@ -615,17 +594,17 @@ def get_site_packages_dirs(python_executable: Optional[str]) -> Tuple[List[str],
     """
 
     if python_executable is None:
-        return [], []
+        return [], [], []
     elif python_executable == sys.executable:
         # Use running Python's package dirs
-        site_packages = pyinfo.getsitepackages()
+        site_packages, sys_path = pyinfo.getsearchdirs()
     else:
         # Use subprocess to get the package directory of given Python
         # executable
-        site_packages = ast.literal_eval(
-            subprocess.check_output([python_executable, pyinfo.__file__, 'getsitepackages'],
+        site_packages, sys_path = ast.literal_eval(
+            subprocess.check_output([python_executable, pyinfo.__file__, 'getsearchdirs'],
             stderr=subprocess.PIPE).decode())
-    return expand_site_packages(site_packages)
+    return expand_site_packages(site_packages) + (sys_path,)
 
 
 def expand_site_packages(site_packages: List[str]) -> Tuple[List[str], List[str]]:
@@ -758,10 +737,8 @@ 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:
+    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
@@ -770,15 +747,10 @@ def compute_search_paths(sources: List[BuildSource],
             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(egg_dirs + site_packages + sys_path),
                        typeshed_path=tuple(lib_path))
 
 

mypy/pyinfo.py

@@ -6,8 +6,10 @@
 library found in Python 2. This file is run each mypy run, so it should be kept as fast as
 possible.
 """
+import os
 import site
 import sys
+import sysconfig
 
 if __name__ == '__main__':
     sys.path = sys.path[1:]  # we don't want to pick up mypy.types
@@ -17,12 +19,27 @@
     from typing import List, Tuple
 
 
-def getprefixes():
-    # type: () -> Tuple[str, str]
-    return sys.base_prefix, sys.prefix
+def getsearchdirs():
+    # type: () -> Tuple[List[str], List[str]]
+    site_packages = _getsitepackages()
 
+    # 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(site_packages + [cwd, stdlib_zip, stdlib, stdlib_ext])
 
-def getsitepackages():
+    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]
     if hasattr(site, 'getusersitepackages') and hasattr(site, 'getsitepackages'):
         user_dir = site.getusersitepackages()
@@ -33,10 +50,8 @@ def getsitepackages():
 
 
 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

@@ -56,7 +56,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'] += ':' + extra_path
     process = subprocess.Popen(fixed + args,
                                stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE,

test-data/unit/cmdline.test

@@ -363,6 +363,30 @@ 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
+[file pypath/a/__init__.py]
+x = 0
+x += ''  # Error reported here
+[file pypath/a/py.typed]
+[out]
+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]