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

Find 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.

Author: Brandon Carpenter

mypy/modulefinder.py

@@ -6,6 +6,8 @@
 import ast
 import collections
 import functools
+import glob
+import json
 import os
 import re
 import subprocess
@@ -752,7 +754,42 @@ def get_search_dirs(python_executable: Optional[str]) -> List[str]:
             raise CompileError(
                 [f"mypy: Invalid python executable '{python_executable}': {reason}"]
             ) from err
-    return sys_path
+    return sys_path + find_editable_dirs(sys_path)
+
+
+def find_editable_dirs(search_paths: List[str]) -> List[str]:
+    """Find directories of editable packages."""
+    editable_paths = []
+    for search_path in search_paths:
+        for meta_path in glob.glob('*.dist-info/direct_url.json', root_dir=search_path):
+            path = _parse_direct_url_file(os.path.join(search_path, meta_path))
+            if path and path not in search_path and path not in editable_paths:
+                editable_paths.append(path)
+    return editable_paths
+
+
+def _parse_direct_url_file(path: str) -> Optional[str]:
+    """Get the path of an editable package using direct_url.json.
+
+    See https://www.python.org/dev/peps/pep-0610/
+    """
+    try:
+        file = open(path, encoding="utf-8")
+    except OSError:
+        return
+    with file:
+        try:
+            direct_url = json.load(file)
+        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:]
+        # Path will already be absolute, except for those in the unit tests
+        return os.path.normpath(path)
 
 
 def add_py2_mypypath_entries(mypypath: List[str]) -> List[str]:

mypy/test/testmodulefinder.py

@@ -5,6 +5,7 @@
     FindModuleCache,
     SearchPaths,
     ModuleNotFoundReason,
+    find_editable_dirs,
 )
 
 from mypy.test.helpers import Suite, assert_equal
@@ -153,12 +154,12 @@ def setUp(self) -> None:
             os.path.join(self.package_dir, "..", "not-a-directory"),
             os.path.join(self.package_dir, "..", "modulefinder-src"),
             self.package_dir,
-        )
+        ) + tuple(find_editable_dirs([self.package_dir]))
 
         self.search_paths = SearchPaths(
             python_path=(),
             mypy_path=(os.path.join(data_path, "pkg1"),),
-            package_path=tuple(package_paths),
+            package_path=package_paths,
             typeshed_path=(),
         )
         options = Options()
@@ -172,6 +173,14 @@ def setUp(self) -> None:
     def path(self, *parts: str) -> str:
         return os.path.join(self.package_dir, *parts)
 
+    def test__find_editable_dirs(self):
+        found_dirs = find_editable_dirs([self.package_dir])
+        found_dirs.sort()
+        expected = [os.path.join(*p)
+                    for p in [("test-data", "packages", "editable_pkg_typed-src"),
+                              ("test-data", "packages", "editable_pkg_untyped-src")]]
+        assert_equal(expected, found_dirs)
+
     def test__packages_with_ns(self) -> None:
         cases = [
             # Namespace package with py.typed
@@ -213,6 +222,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")),
@@ -275,6 +289,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/editable_pkg_typed-src/editable_pkg_typed/__init__.py

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

test-data/packages/editable_pkg_typed-src/editable_pkg_typed/py.typed

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

test-data/packages/editable_pkg_untyped-src/editable_pkg_untyped/__init__.py

@@ -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"
+    }
+}