Fix search paths for -p (python#9387)

huguesb · web-flow · commit d7553fe181f9 · 2020-09-08T15:06:39.000-07:00
`-p` is similar to `-m` except that it will recursively typecheck submodules. Unfortunately the early module walk was being done with search paths that did not benefit from site_packages for the selected python executable, which resulted in some packages being typechecked fine with `-m`, but rejected as `not found` by `-p`. Fixes python#9386
diff --git a/mypy/main.py b/mypy/main.py
@@ -14,7 +14,10 @@
 from mypy import defaults
 from mypy import state
 from mypy import util
-from mypy.modulefinder import BuildSource, FindModuleCache, mypy_path, SearchPaths
+from mypy.modulefinder import (
+    BuildSource, FindModuleCache, SearchPaths,
+    get_site_packages_dirs, mypy_path,
+)
 from mypy.find_sources import create_source_list, InvalidSourceList
 from mypy.fscache import FileSystemCache
 from mypy.errors import CompileError
@@ -921,7 +924,11 @@ def set_strict_flags() -> None:
     # Set target.
     if special_opts.modules + special_opts.packages:
         options.build_type = BuildType.MODULE
-        search_paths = SearchPaths((os.getcwd(),), tuple(mypy_path() + options.mypy_path), (), ())
+        egg_dirs, site_packages = get_site_packages_dirs(options.python_executable)
+        search_paths = SearchPaths((os.getcwd(),),
+                                   tuple(mypy_path() + options.mypy_path),
+                                   tuple(egg_dirs + site_packages),
+                                   ())
         targets = []
         # TODO: use the same cache that the BuildManager will
         cache = FindModuleCache(search_paths, fscache, options, special_opts.packages)
diff --git a/mypy/modulefinder.py b/mypy/modulefinder.py
@@ -473,14 +473,16 @@ def default_lib_path(data_dir: str,
 
 
 @functools.lru_cache(maxsize=None)
-def get_site_packages_dirs(python_executable: str) -> Tuple[List[str], List[str]]:
+def get_site_packages_dirs(python_executable: Optional[str]) -> Tuple[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
     package directories. To avoid repeatedly calling a subprocess (which can be slow!) we
     lru_cache the results."""
 
-    if python_executable == sys.executable:
+    if python_executable is None:
+        return [], []
+    elif python_executable == sys.executable:
         # Use running Python's package dirs
         site_packages = sitepkgs.getsitepackages()
     else:
@@ -598,11 +600,7 @@ def compute_search_paths(sources: List[BuildSource],
     if alt_lib_path:
         mypypath.insert(0, alt_lib_path)
 
-    if options.python_executable is None:
-        egg_dirs = []  # type: List[str]
-        site_packages = []  # type: List[str]
-    else:
-        egg_dirs, site_packages = get_site_packages_dirs(options.python_executable)
+    egg_dirs, site_packages = get_site_packages_dirs(options.python_executable)
     for site_dir in site_packages:
         assert site_dir not in lib_path
         if (site_dir in mypypath or
diff --git a/mypy/test/testpep561.py b/mypy/test/testpep561.py
@@ -120,11 +120,15 @@ def test_pep561(testcase: DataDrivenTestCase) -> None:
             old_dir = os.getcwd()
             os.chdir(venv_dir)
         try:
-            program = testcase.name + '.py'
-            with open(program, 'w', encoding='utf-8') as f:
-                for s in testcase.input:
-                    f.write('{}\n'.format(s))
-            cmd_line = mypy_args + [program, '--no-incremental', '--no-error-summary']
+            cmd_line = list(mypy_args)
+            has_program = not ('-p' in cmd_line or '--package' in cmd_line)
+            if has_program:
+                program = testcase.name + '.py'
+                with open(program, 'w', encoding='utf-8') as f:
+                    for s in testcase.input:
+                        f.write('{}\n'.format(s))
+                cmd_line.append(program)
+            cmd_line.extend(['--no-incremental', '--no-error-summary'])
             if python_executable != sys.executable:
                 cmd_line.append('--python-executable={}'.format(python_executable))
             if testcase.files != []:
@@ -135,7 +139,8 @@ def test_pep561(testcase: DataDrivenTestCase) -> None:
             output = []
             # Type check the module
             out, err, returncode = mypy.api.run(cmd_line)
-            os.remove(program)
+            if has_program:
+                os.remove(program)
             # split lines, remove newlines, and remove directory of test case
             for line in (out + err).splitlines():
                 if line.startswith(test_temp_dir + os.sep):
diff --git a/test-data/unit/pep561.test b/test-data/unit/pep561.test
@@ -20,6 +20,11 @@ reveal_type(a)
 [out]
 testTypedPkgSimple.py:5: note: Revealed type is 'builtins.tuple[builtins.str]'
 
+[case testTypedPkgSimplePackageSearchPath]
+# pkgs: typedpkg
+# flags: -p typedpkg.pkg
+[out]
+
 [case testTypedPkg_config_nositepackages]
 # pkgs: typedpkg
 from typedpkg.sample import ex
