Find Python 2 variants of stub packages

JukkaL · JukkaL · commit b4e1b18e8d0a · 2020-06-29T16:48:34.000+01:00
This is currently needed for six.
diff --git a/mypy/modulefinder.py b/mypy/modulefinder.py
@@ -120,6 +120,7 @@ def __init__(self,
         self.options = options
         self.ns_packages = ns_packages or []  # type: List[str]
         self.stdlib_py_versions = load_stdlib_py_versions()
+        self.python2 = options and options.python_version[0] == 2
 
     def clear(self) -> None:
         self.results.clear()
@@ -230,6 +231,12 @@ def _find_module(self, id: str, use_typeshed: bool) -> ModuleSearchResult:
         for pkg_dir in self.search_paths.package_path:
             stub_name = components[0] + '-stubs'
             stub_dir = os.path.join(pkg_dir, stub_name)
+            if self.python2:
+                alt_stub_name = components[0] + '-python2-stubs'
+                alt_stub_dir = os.path.join(pkg_dir, alt_stub_name)
+                if fscache.isdir(alt_stub_dir):
+                    stub_name = alt_stub_name
+                    stub_dir = alt_stub_dir
             if fscache.isdir(stub_dir) and self._is_compatible_stub_package(stub_dir):
                 stub_typed_file = os.path.join(stub_dir, 'py.typed')
                 stub_components = [stub_name] + components[1:]
@@ -291,7 +298,11 @@ def _find_module(self, id: str, use_typeshed: bool) -> ModuleSearchResult:
             # 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
+                suffix = '-stubs'
+                if self.python2:
+                    if os.path.isdir(base_path + '-python2-stubs'):
+                        suffix = '-python2-stubs'
+                path_stubs = base_path + suffix + sepinit + extension
                 if fscache.isfile_case(path, dir_prefix):
                     has_init = True
                     if verify and not verify_module(fscache, id, path, dir_prefix):
@@ -367,12 +378,12 @@ def _is_compatible_stub_package(self, stub_dir: str) -> bool:
         whether the stubs are compatible with Python 2 and 3.
         """
         metadata_fnam = os.path.join(stub_dir, 'METADATA.toml')
-        if os.path.isfile(metadata_fnam) and self.options:
+        if os.path.isfile(metadata_fnam):
             # Delay import for a possible minor performance win.
             import toml
             with open(metadata_fnam, 'r') as f:
                 metadata = toml.load(f)
-            if self.options.python_version[0] == 2:
+            if self.python2:
                 return bool(metadata.get('python2', False))
             else:
                 return bool(metadata.get('python3', True))
