Add support for PEP-610 editable packages (python#12313)

Discover paths to editable packages by looking at direct_url.json files
-- installed by pip when the -e/--editable flag is used -- and add to
the package search paths.

mypy/modulefinder.py

@@ -6,6 +6,7 @@
 import ast
 import collections
 import functools
+import json
 import os
 import re
 import subprocess
@@ -657,13 +658,41 @@ def expand_site_packages(site_packages: List[str]) -> Tuple[List[str], List[str]
     for dir in site_packages:
         if not os.path.isdir(dir):
             continue
-        pth_filenames = sorted(name for name in os.listdir(dir) if name.endswith(".pth"))
-        for pth_filename in pth_filenames:
+        dir_listing = sorted(os.listdir(dir))
+        for dist_dir in [name for name in dir_listing
+                         if name.endswith(".dist-info")]:
+            egg_dirs.extend(_parse_direct_url_file(dir, dist_dir))
+        for pth_filename in [name for name in dir_listing
+                             if name.endswith(".pth")]:
             egg_dirs.extend(_parse_pth_file(dir, pth_filename))
 
     return egg_dirs, site_packages
 
 
+def _parse_direct_url_file(dir: str, dist_dir: str) -> Iterator[str]:
+    """Get the path of an editable package using direct_url.json.
+
+    See https://www.python.org/dev/peps/pep-0610/
+    """
+    url_file = os.path.join(dir, dist_dir, "direct_url.json")
+    try:
+        f = open(url_file, "r", encoding="utf-8")
+    except OSError:
+        return
+    with f:
+        try:
+            direct_url = json.load(f)
+        except (ValueError, IOError):
+            return
+    try:
+        url = str(direct_url["url"]) if direct_url["dir_info"]["editable"] else ""
+    except (LookupError, TypeError):
+        return
+    if url.startswith("file://"):
+        path = url[7:]
+        yield path
+
+
 def _parse_pth_file(dir: str, pth_filename: str) -> Iterator[str]:
     """
     Mimics a subset of .pth import hook from Lib/site.py

mypy/test/testmodulefinder.py

@@ -209,6 +209,11 @@ def test__packages_with_ns(self) -> None:
             ("standalone", ModuleNotFoundReason.FOUND_WITHOUT_TYPE_HINTS),
             ("standalone.standalone_var", ModuleNotFoundReason.FOUND_WITHOUT_TYPE_HINTS),
 
+            # Packages found by following direct_url.json files
+            ("editable_pkg_typed", os.path.join("test-data", "packages", "editable_pkg_typed-src",
+                                                "editable_pkg_typed", "__init__.py")),
+            ("editable_pkg_untyped", ModuleNotFoundReason.FOUND_WITHOUT_TYPE_HINTS),
+
             # Packages found by following .pth files
             ("baz_pkg", self.path("baz", "baz_pkg", "__init__.py")),
             ("ns_baz_pkg.a", self.path("baz", "ns_baz_pkg", "a.py")),
@@ -271,6 +276,11 @@ def test__packages_without_ns(self) -> None:
             ("standalone", ModuleNotFoundReason.FOUND_WITHOUT_TYPE_HINTS),
             ("standalone.standalone_var", ModuleNotFoundReason.FOUND_WITHOUT_TYPE_HINTS),
 
+            # Packages found by following direct_url.json files
+            ("editable_pkg_typed", os.path.join("test-data", "packages", "editable_pkg_typed-src",
+                                                "editable_pkg_typed", "__init__.py")),
+            ("editable_pkg_untyped", ModuleNotFoundReason.FOUND_WITHOUT_TYPE_HINTS),
+
             # Packages found by following .pth files
             ("baz_pkg", self.path("baz", "baz_pkg", "__init__.py")),
             ("ns_baz_pkg.a", ModuleNotFoundReason.NOT_FOUND),

test-data/packages/modulefinder-site-packages/editable_pkg_typed-0.dist-info/direct_url.json

@@ -0,0 +1,6 @@
+{
+    "url": "file://test-data/packages/editable_pkg_typed-src",
+    "dir_info": {
+        "editable": true
+    }
+}

test-data/packages/modulefinder-site-packages/editable_pkg_untyped-0.dist-info/direct_url.json

@@ -0,0 +1,6 @@
+{
+    "url": "file://test-data/packages/editable_pkg_untyped-src",
+    "dir_info": {
+        "editable": true
+    }
+}

test-data/packages/modulefinder-site-packages/foo-0.dist-info/direct_url.json

@@ -0,0 +1,7 @@
+{
+    "comment": "This file is in place only to ensure editable features don't break on non-editable packages.",
+    "url": "https://packages.example.com/foo/0/foo-0.tar.gz",
+    "archive_info": {
+        "hash": "sha256=2dc6b5a470a1bde68946f263f1af1515a2574a150a30d6ce02c6ff742fcc0db8"
+    }
+}