Support egg/setuptools packages for PEP 561 searching. (#5282)

This adds support for setuptool's egg format, which includes support for editable installs (`pip install -e .`/`python setup.py develop`)!

Setuptools creates its own directory to put the package we want to find. These directories are listed in `easy-install.pth` in a site-package directory.

Fixes #5007.

Author: emmatyping

docs/source/installed_packages.rst

@@ -43,7 +43,7 @@ If you would like to publish a library package to a package repository (e.g.
 PyPI) for either internal or external use in type checking, packages that
 supply type information via type comments or annotations in the code should put
 a ``py.typed`` in their package directory. For example, with a directory
-structure as follows:
+structure as follows
 
 .. code-block:: text
 
@@ -53,7 +53,7 @@ structure as follows:
         lib.py
         py.typed
 
-the setup.py might look like:
+the setup.py might look like
 
 .. code-block:: python
 
@@ -67,8 +67,13 @@ the setup.py might look like:
         packages=["package_a"]
     )
 
+.. note::
+
+   If you use setuptools, you must pass the option ``zip_safe=False`` to
+   ``setup()``, or mypy will not be able to find the installed package.
+
 Some packages have a mix of stub files and runtime files. These packages also
-require a ``py.typed`` file. An example can be seen below:
+require a ``py.typed`` file. An example can be seen below
 
 .. code-block:: text
 

mypy/build.py

@@ -199,26 +199,43 @@ def default_flush_errors(new_messages: List[str], is_serious: bool) -> None:
 
 
 @functools.lru_cache(maxsize=None)
-def _get_site_packages_dirs(python_executable: Optional[str]) -> List[str]:
+def _get_site_packages_dirs(python_executable: Optional[str],
+                            fscache: FileSystemCache) -> List[str]:
     """Find package directories for given python.
 
     This runs a subprocess call, which generates a list of the site package directories.
     To avoid repeatedly calling a subprocess (which can be slow!) we lru_cache the results."""
+    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))
+
     if python_executable is None:
         return []
     if python_executable == sys.executable:
         # Use running Python's package dirs
-        return sitepkgs.getsitepackages()
+        site_packages = sitepkgs.getsitepackages()
     else:
         # Use subprocess to get the package directory of given Python
         # executable
-        return ast.literal_eval(subprocess.check_output([python_executable, sitepkgs.__file__],
-                                stderr=subprocess.PIPE).decode())
+        site_packages = ast.literal_eval(
+            subprocess.check_output([python_executable, sitepkgs.__file__],
+            stderr=subprocess.PIPE).decode())
+    egg_dirs = []
+    for dir in site_packages:
+        pth = os.path.join(dir, 'easy-install.pth')
+        if fscache.isfile(pth):
+            with open(pth) as f:
+                egg_dirs.extend([make_abspath(d.rstrip(), dir) for d in f.readlines()])
+    return egg_dirs + site_packages
 
 
 def compute_search_paths(sources: List[BuildSource],
                      options: Options,
                      data_dir: str,
+                     fscache: FileSystemCache,
                      alt_lib_path: Optional[str] = None) -> SearchPaths:
     """Compute the search paths as specified in PEP 561.
 
@@ -275,7 +292,7 @@ def compute_search_paths(sources: List[BuildSource],
     if alt_lib_path:
         mypypath.insert(0, alt_lib_path)
 
-    package_path = tuple(_get_site_packages_dirs(options.python_executable))
+    package_path = tuple(_get_site_packages_dirs(options.python_executable, fscache))
     for site_dir in package_path:
         assert site_dir not in lib_path
         if site_dir in mypypath:
@@ -306,7 +323,7 @@ def _build(sources: List[BuildSource],
     data_dir = default_data_dir(bin_dir)
     fscache = fscache or FileSystemCache()
 
-    search_paths = compute_search_paths(sources, options, data_dir, alt_lib_path)
+    search_paths = compute_search_paths(sources, options, data_dir, fscache, alt_lib_path)
 
     reports = Reports(data_dir, options.report_dirs)
     source_set = BuildSourceSet(sources)

mypy/dmypy_server.py

@@ -327,7 +327,7 @@ def fine_grained_increment(self, sources: List[mypy.build.BuildSource]) -> Dict[
         self.update_sources(sources)
         changed, removed = self.find_changed(sources)
         manager.search_paths = mypy.build.compute_search_paths(
-            sources, manager.options, manager.data_dir)
+            sources, manager.options, manager.data_dir, mypy.build.FileSystemCache())
         t1 = time.time()
         messages = self.fine_grained_manager.update(changed, removed)
         t2 = time.time()

mypy/test/testpep561.py

@@ -6,7 +6,7 @@
 from unittest import TestCase, main
 
 import mypy.api
-from mypy.build import _get_site_packages_dirs
+from mypy.build import _get_site_packages_dirs, FileSystemCache
 from mypy.test.config import package_path
 from mypy.test.helpers import run_command
 from mypy.util import try_find_python2_interpreter
@@ -67,10 +67,22 @@ def virtualenv(self,
                 yield venv_dir, os.path.abspath(os.path.join(venv_dir, 'bin', 'python'))
 
     def install_package(self, pkg: str,
-                        python_executable: str = sys.executable) -> None:
+                        python_executable: str = sys.executable,
+                        use_pip: bool = True,
+                        editable: bool = False) -> None:
         """Context manager to temporarily install a package from test-data/packages/pkg/"""
         working_dir = os.path.join(package_path, pkg)
-        install_cmd = [python_executable, '-m', 'pip', 'install', '.']
+        if use_pip:
+            install_cmd = [python_executable, '-m', 'pip', 'install']
+            if editable:
+                install_cmd.append('-e')
+            install_cmd.append('.')
+        else:
+            install_cmd = [python_executable, 'setup.py']
+            if editable:
+                install_cmd.append('develop')
+            else:
+                install_cmd.append('install')
         returncode, lines = run_command(install_cmd, cwd=working_dir)
         if returncode != 0:
             self.fail('\n'.join(lines))
@@ -92,7 +104,7 @@ def tearDown(self) -> None:
 
     def test_get_pkg_dirs(self) -> None:
         """Check that get_package_dirs works."""
-        dirs = _get_site_packages_dirs(sys.executable)
+        dirs = _get_site_packages_dirs(sys.executable, FileSystemCache())
         assert dirs
 
     def test_typedpkg_stub_package(self) -> None:
@@ -155,6 +167,28 @@ def test_typedpkg_python2(self) -> None:
                     venv_dir=venv_dir,
                 )
 
+    def test_typedpkg_egg(self) -> None:
+        with self.virtualenv() as venv:
+            venv_dir, python_executable = venv
+            self.install_package('typedpkg', python_executable, use_pip=False)
+            check_mypy_run(
+                [self.tempfile],
+                python_executable,
+                expected_out=self.msg_tuple,
+                venv_dir=venv_dir,
+            )
+
+    def test_typedpkg_editable(self) -> None:
+        with self.virtualenv() as venv:
+            venv_dir, python_executable = venv
+            self.install_package('typedpkg', python_executable, editable=True)
+            check_mypy_run(
+                [self.tempfile],
+                python_executable,
+                expected_out=self.msg_tuple,
+                venv_dir=venv_dir,
+            )
+
 
 if __name__ == '__main__':
     main()

test-data/packages/typedpkg/setup.py

@@ -2,7 +2,7 @@
 This setup file installs packages to test mypy's PEP 561 implementation
 """
 
-from distutils.core import setup
+from setuptools import setup
 
 setup(
     name='typedpkg',
@@ -11,4 +11,5 @@
     package_data={'typedpkg': ['py.typed']},
     packages=['typedpkg'],
     include_package_data=True,
+    zip_safe=False,
 )

test-requirements.txt

@@ -10,3 +10,4 @@ typed-ast>=1.1.0,<1.2.0
 typing>=3.5.2; python_version < '3.5'
 py>=1.5.2
 virtualenv
+setuptools