Add support for partial stub packages (#5227)

This should conform to PEP 561 as laid out in
https://www.python.org/dev/peps/pep-0561/#partial-stub-packages.

Author: emmatyping

mypy/build.py

@@ -862,6 +862,10 @@ def _get_site_packages_dirs(python_executable: Optional[str]) -> List[str]:
                                 stderr=subprocess.PIPE).decode())
 
 
+# Search paths are a two-tuple of path and whether to verify the module
+SearchPaths = List[Tuple[str, bool]]
+
+
 class FindModuleCache:
     """Module finder with integrated cache.
 
@@ -875,16 +879,16 @@ class FindModuleCache:
 
     def __init__(self, fscache: Optional[FileSystemCache] = None) -> None:
         self.fscache = fscache or FileSystemCache()
-        # Cache find_lib_path_dirs: (dir_chain, lib_path)
-        self.dirs = {}  # type: Dict[Tuple[str, Tuple[str, ...]], List[str]]
+        # 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]]
 
     def clear(self) -> None:
         self.results.clear()
         self.dirs.clear()
 
-    def find_lib_path_dirs(self, dir_chain: str, lib_path: Tuple[str, ...]) -> List[str]:
+    def find_lib_path_dirs(self, dir_chain: str, lib_path: Tuple[str, ...]) -> SearchPaths:
         # 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
@@ -894,13 +898,13 @@ def find_lib_path_dirs(self, dir_chain: str, lib_path: Tuple[str, ...]) -> List[
             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, ...]) -> List[str]:
+    def _find_lib_path_dirs(self, dir_chain: str, lib_path: Tuple[str, ...]) -> SearchPaths:
         dirs = []
         for pathitem in lib_path:
             # e.g., '/usr/lib/python3.4/foo/bar'
             dir = os.path.normpath(os.path.join(pathitem, dir_chain))
             if self.fscache.isdir(dir):
-                dirs.append(dir)
+                dirs.append((dir, True))
         return dirs
 
     def find_module(self, id: str, lib_path: Tuple[str, ...],
@@ -933,13 +937,26 @@ def _find_module(self, id: str, lib_path: Tuple[str, ...],
             typed_file = os.path.join(pkg_dir, components[0], 'py.typed')
             stub_dir = os.path.join(pkg_dir, stub_name)
             if fscache.isdir(stub_dir):
+                stub_typed_file = os.path.join(stub_dir, 'py.typed')
                 stub_components = [stub_name] + components[1:]
                 path = os.path.join(pkg_dir, *stub_components[:-1])
                 if fscache.isdir(path):
-                    third_party_stubs_dirs.append(path)
+                    if fscache.isfile(stub_typed_file):
+                        # Stub packages can have a py.typed file, which must include
+                        # 'partial\n' to make the package partial
+                        # Partial here means that mypy should look at the runtime
+                        # package if installed.
+                        if fscache.read(stub_typed_file).decode().strip() == 'partial':
+                            runtime_path = os.path.join(pkg_dir, dir_chain)
+                        third_party_inline_dirs.append((runtime_path, True))
+                        # if the package is partial, we don't verify the module, as
+                        # the partial stub package may not have a __init__.pyi
+                        third_party_stubs_dirs.append((path, False))
+                    else:
+                        third_party_stubs_dirs.append((path, True))
             elif fscache.isfile(typed_file):
                 path = os.path.join(pkg_dir, dir_chain)
-                third_party_inline_dirs.append(path)
+                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
 
@@ -949,20 +966,26 @@ def _find_module(self, id: str, lib_path: Tuple[str, ...],
         # Now just look for 'baz.pyi', 'baz/__init__.py', etc., inside those directories.
         seplast = os.sep + components[-1]  # so e.g. '/baz'
         sepinit = os.sep + '__init__'
-        for base_dir in candidate_base_dirs:
+        for base_dir, verify in candidate_base_dirs:
             base_path = base_dir + seplast  # so e.g. '/usr/lib/python3.4/foo/bar/baz'
             # Prefer package over module, i.e. baz/__init__.py* over baz.py*.
             for extension in PYTHON_EXTENSIONS:
                 path = base_path + sepinit + extension
                 path_stubs = base_path + '-stubs' + sepinit + extension
-                if fscache.isfile_case(path) and verify_module(fscache, id, path):
+                if fscache.isfile_case(path):
+                    if verify and not verify_module(fscache, id, path):
+                        continue
                     return path
-                elif fscache.isfile_case(path_stubs) and verify_module(fscache, id, path_stubs):
+                elif fscache.isfile_case(path_stubs):
+                    if verify and not verify_module(fscache, id, path_stubs):
+                        continue
                     return path_stubs
             # No package, look for module.
             for extension in PYTHON_EXTENSIONS:
                 path = base_path + extension
-                if fscache.isfile_case(path) and verify_module(fscache, id, path):
+                if fscache.isfile_case(path):
+                    if verify and not verify_module(fscache, id, path):
+                        continue
                     return path
         return None
 

mypy/test/testpep561.py

@@ -13,6 +13,7 @@
 
 SIMPLE_PROGRAM = """
 from typedpkg.sample import ex
+from typedpkg import dne
 a = ex([''])
 reveal_type(a)
 """
@@ -79,10 +80,12 @@ def setUp(self) -> None:
         self.tempfile = os.path.join(self.temp_file_dir.name, 'simple.py')
         with open(self.tempfile, 'w+') as file:
             file.write(SIMPLE_PROGRAM)
+        self.msg_dne = \
+            "{}:3: error: Module 'typedpkg' has no attribute 'dne'\n".format(self.tempfile)
         self.msg_list = \
-            "{}:4: error: Revealed type is 'builtins.list[builtins.str]'\n".format(self.tempfile)
+            "{}:5: error: Revealed type is 'builtins.list[builtins.str]'\n".format(self.tempfile)
         self.msg_tuple = \
-            "{}:4: error: Revealed type is 'builtins.tuple[builtins.str]'\n".format(self.tempfile)
+            "{}:5: error: Revealed type is 'builtins.tuple[builtins.str]'\n".format(self.tempfile)
 
     def tearDown(self) -> None:
         self.temp_file_dir.cleanup()
@@ -99,7 +102,7 @@ def test_typedpkg_stub_package(self) -> None:
             check_mypy_run(
                 [self.tempfile],
                 python_executable,
-                expected_out=self.msg_list,
+                expected_out=self.msg_dne + self.msg_list,
                 venv_dir=venv_dir,
             )
 
@@ -135,7 +138,7 @@ def test_typedpkg_stubs_python2(self) -> None:
                 check_mypy_run(
                     [self.tempfile],
                     py2,
-                    expected_out=self.msg_list,
+                    expected_out=self.msg_dne + self.msg_list,
                     venv_dir=venv_dir,
                 )
 

test-data/packages/typedpkg-stubs/setup.py

@@ -8,6 +8,6 @@
     name='typedpkg-stubs',
     author="The mypy team",
     version='0.1',
-    package_data={'typedpkg-stubs': ['sample.pyi', '__init__.pyi']},
+    package_data={'typedpkg-stubs': ['sample.pyi', '__init__.pyi', 'py.typed']},
     packages=['typedpkg-stubs'],
 )

test-data/packages/typedpkg-stubs/typedpkg-stubs/py.typed

@@ -0,0 +1 @@
+partial