Implement PEP561

docs/source/command_line.rst

@@ -363,11 +363,27 @@ Here are some more useful flags:
   updates the cache, but regular incremental mode ignores cache files
   written by quick mode.
 
+- ``--python-executable EXECUTABLE`` will have mypy collect type information
+  from `PEP 561`_ compliant packages installed for the Python executable
+  ``EXECUTABLE``. If not provided, mypy will use PEP 561 compliant packages
+  installed for the Python executable running mypy. See
+  :ref:`installed-packages` for more on making PEP 561 compliant packages.
+  This flag will attempt to set ``--python-version`` if not already set.
+
 - ``--python-version X.Y`` will make mypy typecheck your code as if it were
   run under Python version X.Y. Without this option, mypy will default to using
   whatever version of Python is running mypy. Note that the ``-2`` and
   ``--py2`` flags are aliases for ``--python-version 2.7``. See
-  :ref:`version_and_platform_checks` for more about this feature.
+  :ref:`version_and_platform_checks` for more about this feature. This flag
+  will attempt to find a Python executable of the corresponding version to
+  search for `PEP 561`_ compliant packages. If you'd like to disable this,
+  see ``--no-site-packages`` below.
+
+- ``--no-site-packages`` will disable searching for `PEP 561`_ compliant
+  packages. This will also disable searching for a usable Python executable.
+  Use this  flag if mypy cannot find a Python executable for the version of
+  Python being checked, and you don't need to use PEP 561 typed packages.
+  Otherwise, use ``--python-executable``.
 
 - ``--platform PLATFORM`` will make mypy typecheck your code as if it were
   run under the the given operating system. Without this option, mypy will
@@ -450,6 +466,8 @@ For the remaining flags you can read the full ``mypy -h`` output.
 
    Command line flags are liable to change between releases.
 
+.. _PEP 561: https://www.python.org/dev/peps/pep-0561/
+
 .. _integrating-mypy:
 
 Integrating mypy into another Python application

docs/source/index.rst

@@ -30,6 +30,7 @@ Mypy is a static type checker for Python.
    command_line
    config_file
    python36
+   installed_packages
    faq
    cheat_sheet
    cheat_sheet_py3

docs/source/installed_packages.rst

@@ -0,0 +1,114 @@
+.. _installed-packages:
+
+Using Installed Packages
+========================
+
+`PEP 561 <https://www.python.org/dev/peps/pep-0561/>`_ specifies how to mark
+a package as supporting type checking. Below is a summary of how to create
+PEP 561 compatible packages and have mypy use them in type checking.
+
+Making PEP 561 compatible packages
+**********************************
+
+Packages that must be imported at runtime and supply type information should
+put a ``py.typed`` in their package directory. For example, with a directory
+structure as follows:
+
+.. code-block:: text
+
+    setup.py
+    package_a/
+        __init__.py
+        lib.py
+        py.typed
+
+the setup.py might look like:
+
+.. code-block:: python
+
+    from distutils.core import setup
+
+    setup(
+        name="SuperPackageA",
+        author="Me",
+        version="0.1",
+        package_data={"package_a": ["py.typed"]},
+        packages=["package_a"]
+    )
+
+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:
+
+.. code-block:: text
+
+    setup.py
+    package_b/
+        __init__.py
+        lib.py
+        lib.pyi
+        py.typed
+
+the setup.py might look like:
+
+.. code-block:: python
+
+    from distutils.core import setup
+
+    setup(
+        name="SuperPackageB",
+        author="Me",
+        version="0.1",
+        package_data={"package_b": ["py.typed", "lib.pyi"]},
+        packages=["package_b"]
+    )
+
+In this example, both ``lib.py`` and ``lib.pyi`` exist. At runtime, ``lib.py``
+will be used, however mypy will use ``lib.pyi``.
+
+If the package is stub-only (not imported at runtime), the package should have
+a prefix of the runtime package name and a suffix of ``-stubs``.
+A ``py.typed`` file is not needed for stub-only packages. For example, if we
+had stubs for ``package_c``, we might do the following:
+
+.. code-block:: text
+
+    setup.py
+    package_c-stubs/
+        __init__.pyi
+        lib.pyi
+
+the setup.py might look like:
+
+.. code-block:: python
+
+    from distutils.core import setup
+
+    setup(
+        name="SuperPackageC",
+        author="Me",
+        version="0.1",
+        package_data={"package_c-stubs": ["__init__.pyi", "lib.pyi"]},
+        packages=["package_c-stubs"]
+    )
+
+Using PEP 561 compatible packages with mypy
+*******************************************
+
+Generally, you do not need to do anything to use installed packages for the
+Python executable used to run mypy. They should be automatically picked up by
+mypy and used for type checking.
+
+By default, mypy searches for packages installed for the Python executable
+running mypy. It is highly unlikely you want this situation if you have
+installed typed packages in another Python's package directory.
+
+Generally, you can use the ``--python-version`` flag and mypy will try to find
+the correct package directory. If that fails, you can use the
+``--python-executable`` flag to point to the exact executable, and mypy will
+find packages installed for that Python executable.
+
+Note that mypy does not support some more advanced import features, such as zip
+imports, namespace packages, and custom import hooks.
+
+If you do not want to use typed packages, use the ``--no-site-packages`` flag
+to disable searching.

mypy/build.py

@@ -10,17 +10,20 @@
 """
 # TODO: More consistent terminology, e.g. path/fnam, module/id, state/file
 
+import ast
 import binascii
 import collections
 import contextlib
 from distutils.sysconfig import get_python_lib
+import functools
 import gc
 import hashlib
 import json
 import os.path
 import re
 import site
 import stat
+import subprocess
 import sys
 import time
 from os.path import dirname, basename
@@ -696,7 +699,8 @@ def correct_rel_imp(imp: Union[ImportFrom, ImportAll]) -> str:
 
     def is_module(self, id: str) -> bool:
         """Is there a file in the file system corresponding to module id?"""
-        return self.find_module_cache.find_module(id, self.lib_path) is not None
+        return self.find_module_cache.find_module(id, self.lib_path,
+                                                  self.options.python_executable) is not None
 
     def parse_file(self, id: str, path: str, source: str, ignore_errors: bool) -> MypyFile:
         """Parse the source of a file with the given name.
@@ -809,6 +813,44 @@ def remove_cwd_prefix_from_path(fscache: FileSystemCache, p: str) -> str:
     return p
 
 
+USER_SITE_PACKAGES = \
+    'from __future__ import print_function; import site; print(repr(site.getsitepackages()' \
+    ' + [site.getusersitepackages()]))'
+VIRTUALENV_SITE_PACKAGES = \
+    'from distutils.sysconfig import get_python_lib; print(repr([get_python_lib()]))'
+
+
+def call_python(python_executable: str, command: str) -> str:
+    return subprocess.check_output([python_executable, '-c', command],
+                                   stderr=subprocess.PIPE).decode()
+
+
+@functools.lru_cache(maxsize=None)
+def _get_site_packages_dirs(python_executable: Optional[str]) -> List[str]:
+    """Find package directories for given python."""
+    if python_executable is None:
+        return []
+    if python_executable == sys.executable:
+        # Use running Python's package dirs
+        if hasattr(site, 'getusersitepackages') and hasattr(site, 'getsitepackages'):
+            user_dir = site.getusersitepackages()
+            return site.getsitepackages() + [user_dir]
+        # If site doesn't have get(user)sitepackages, we are running in a
+        # virtualenv, and should fall back to get_python_lib
+        return [get_python_lib()]
+    else:
+        # Use subprocess to get the package directory of given Python
+        # executable
+        try:
+            output = call_python(python_executable, USER_SITE_PACKAGES)
+        except subprocess.CalledProcessError:
+            # if no paths are found (raising a CalledProcessError), we fall back on sysconfig,
+            # the python executable is likely in a virtual environment, thus lacking
+            # the needed site methods
+            output = call_python(python_executable, VIRTUALENV_SITE_PACKAGES)
+    return ast.literal_eval(output)
+
+
 class FindModuleCache:
     """Module finder with integrated cache.
 
@@ -822,20 +864,30 @@ class FindModuleCache:
 
     def __init__(self, fscache: Optional[FileSystemMetaCache] = None) -> None:
         self.fscache = fscache or FileSystemMetaCache()
-        # Cache find_module: (id, lib_path) -> result.
-        self.results = {}  # type: Dict[Tuple[str, Tuple[str, ...]], Optional[str]]
+        self.find_lib_path_dirs = functools.lru_cache(maxsize=None)(self._find_lib_path_dirs)
+        self.find_module = functools.lru_cache(maxsize=None)(self._find_module)
 
+    def clear(self) -> None:
+        self.find_module.cache_clear()
+        self.find_lib_path_dirs.cache_clear()
+
+    def _find_lib_path_dirs(self, dir_chain: str, lib_path: Tuple[str, ...],
+                            python_executable: str) -> List[str]:
         # 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
         # in the last component.
-        self.dirs = {}  # type: Dict[Tuple[str, Tuple[str, ...]], List[str]]
-
-    def clear(self) -> None:
-        self.results.clear()
-        self.dirs.clear()
-
-    def _find_module(self, id: str, lib_path: Tuple[str, ...]) -> Optional[str]:
+        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)
+        return dirs
+
+    def _find_module(self, id: str, lib_path: Tuple[str, ...],
+                     python_executable: Optional[str]) -> Optional[str]:
+        """Return the path of the module source file, or None if not found."""
         fscache = self.fscache
 
         # If we're looking for a module like 'foo.bar.baz', it's likely that most of the
@@ -844,15 +896,24 @@ def _find_module(self, id: str, lib_path: Tuple[str, ...]) -> Optional[str]:
         # that will require the same subdirectory.
         components = id.split('.')
         dir_chain = os.sep.join(components[:-1])  # e.g., 'foo/bar'
-        if (dir_chain, lib_path) not in self.dirs:
-            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 fscache.isdir(dir):
-                    dirs.append(dir)
-            self.dirs[dir_chain, lib_path] = dirs
-        candidate_base_dirs = self.dirs[dir_chain, lib_path]
+        # TODO (ethanhs): refactor each path search to its own method with lru_cache
+
+        third_party_dirs = []
+        # Third-party stub/typed packages
+        for pkg_dir in _get_site_packages_dirs(python_executable):
+            stub_name = components[0] + '-stubs'
+            typed_file = os.path.join(pkg_dir, components[0], 'py.typed')
+            stub_dir = os.path.join(pkg_dir, stub_name)
+            if os.path.isdir(stub_dir):
+                stub_components = [stub_name] + components[1:]
+                path = os.path.join(pkg_dir, *stub_components[:-1])
+                if os.path.isdir(path):
+                    third_party_dirs.append(path)
+            elif os.path.isfile(typed_file):
+                path = os.path.join(pkg_dir, dir_chain)
+                third_party_dirs.append(path)
+        candidate_base_dirs = self.find_lib_path_dirs(dir_chain, lib_path,
+                                                      python_executable) + third_party_dirs
 
         # If we're looking for a module like 'foo.bar.baz', then candidate_base_dirs now
         # contains just the subdirectories 'foo/bar' that actually exist under the
@@ -865,26 +926,21 @@ def _find_module(self, id: str, lib_path: Tuple[str, ...]) -> Optional[str]:
             # 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):
                     return path
+                elif fscache.isfile_case(path_stubs) and verify_module(fscache, id, path_stubs):
+                    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):
                     return path
         return None
 
-    def find_module(self, id: str, lib_path_arg: Iterable[str]) -> Optional[str]:
-        """Return the path of the module source file, or None if not found."""
-        lib_path = tuple(lib_path_arg)
-
-        key = (id, lib_path)
-        if key not in self.results:
-            self.results[key] = self._find_module(id, lib_path)
-        return self.results[key]
-
-    def find_modules_recursive(self, module: str, lib_path: List[str]) -> List[BuildSource]:
-        module_path = self.find_module(module, lib_path)
+    def find_modules_recursive(self, module: str, lib_path: Tuple[str, ...],
+                               python_executable: Optional[str]) -> List[BuildSource]:
+        module_path = self.find_module(module, lib_path, python_executable)
         if not module_path:
             return []
         result = [BuildSource(module_path, module, None)]
@@ -904,13 +960,15 @@ def find_modules_recursive(self, module: str, lib_path: List[str]) -> List[Build
                         (os.path.isfile(os.path.join(abs_path, '__init__.py')) or
                         os.path.isfile(os.path.join(abs_path, '__init__.pyi'))):
                     hits.add(item)
-                    result += self.find_modules_recursive(module + '.' + item, lib_path)
+                    result += self.find_modules_recursive(module + '.' + item, lib_path,
+                                                          python_executable)
                 elif item != '__init__.py' and item != '__init__.pyi' and \
                         item.endswith(('.py', '.pyi')):
                     mod = item.split('.')[0]
                     if mod not in hits:
                         hits.add(mod)
-                        result += self.find_modules_recursive(module + '.' + mod, lib_path)
+                        result += self.find_modules_recursive(module + '.' + mod, lib_path,
+                                                              python_executable)
         return result
 
 
@@ -1523,7 +1581,8 @@ def __init__(self,
                 # difference and just assume 'builtins' everywhere,
                 # which simplifies code.
                 file_id = '__builtin__'
-            path = manager.find_module_cache.find_module(file_id, manager.lib_path)
+            path = manager.find_module_cache.find_module(file_id, manager.lib_path,
+                                                         manager.options.python_executable)
             if path:
                 # For non-stubs, look at options.follow_imports:
                 # - normal (default) -> fully analyze
@@ -2037,8 +2096,6 @@ def dispatch(sources: List[BuildSource], manager: BuildManager) -> Graph:
                       stubs_found=sum(g.path is not None and g.path.endswith('.pyi')
                                       for g in graph.values()),
                       graph_load_time=(t1 - t0),
-                      fm_cache_size=len(manager.find_module_cache.results),
-                      fm_dir_cache_size=len(manager.find_module_cache.dirs),
                       )
     if not graph:
         print("Nothing to do?!")