fix version ranges (#16)

JelleZijlstra · web-flow · commit 8eb1cc9d74fa · 2021-05-05T21:15:30.000-07:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,13 +1,16 @@
 [tool.black]
-target_version = ['py36', 'py37', 'py38']
-include = '\.py$'
+target_version = ['py36', 'py37', 'py38', 'py39']
+include = '\.pyi?$'
+skip-magic-trailing-comma = true
+experimental-string-processing = true
 
 exclude = '''
 /(
     \.git
   | \.mypy_cache
   | \.tox
   | \.venv
+  | typeshed_client/typeshed
 )/
 '''
 
diff --git a/tests/test.py b/tests/test.py
@@ -1,4 +1,5 @@
 from pathlib import Path
+import sys
 from typed_ast import ast3
 import typeshed_client
 from typeshed_client.finder import (
@@ -45,6 +46,8 @@ def test_get_stub_file(self) -> None:
 
         self.check("subdir", (3, 6), TEST_TYPESHED / "subdir/__init__.pyi")
         self.check("subdir.overloads", (3, 7), TEST_TYPESHED / "subdir/overloads.pyi")
+        self.check("subdir", (2, 7), TEST_TYPESHED / "@python2/subdir.pyi")
+        self.check("subdir.overloads", (2, 7), None)
 
     def test_third_party(self) -> None:
         self.check("thirdparty", (3, 6), PACKAGES / "thirdparty-stubs/__init__.pyi")
@@ -57,8 +60,7 @@ def test_get_all_stub_files(self) -> None:
             {
                 ("thirdparty", PACKAGES / "thirdparty-stubs/__init__.pyi"),
                 ("nostubs", PACKAGES / "nostubs/__init__.pyi"),
-                ("subdir", TEST_TYPESHED / "subdir/__init__.pyi"),
-                ("subdir.overloads", TEST_TYPESHED / "subdir/overloads.pyi"),
+                ("subdir", TEST_TYPESHED / "@python2/subdir.pyi"),
                 ("py2only", TEST_TYPESHED / "@python2/py2only.pyi"),
                 ("lib", TEST_TYPESHED / "@python2/lib.pyi"),
                 ("conditions", TEST_TYPESHED / "conditions.pyi"),
diff --git a/tests/typeshed/@python2/subdir.pyi b/tests/typeshed/@python2/subdir.pyi
@@ -0,0 +1 @@
+x: int
diff --git a/tests/typeshed/overloads.pyi b/tests/typeshed/overloads.pyi
@@ -1,4 +1,5 @@
 from typing import overload
+
 @overload
 def overloaded(x: int) -> None: ...
 @overload
diff --git a/tox.ini b/tox.ini
@@ -5,7 +5,7 @@ envlist = py36,py37,py38,py39,black
 [testenv]
 deps =
     typed_ast >= 1.0.3
-    black == 20.8b1
+    black == 21.5b0
 commands =
     python tests/test.py
 
diff --git a/typeshed_client/finder.py b/typeshed_client/finder.py
@@ -137,7 +137,12 @@ def get_all_stub_files(
                 module_name = entry.name[: -len(".pyi")]
             else:
                 continue
-            if search_context.version < versions[module_name]:
+            version = versions[module_name]
+            if search_context.version < version.min:
+                continue
+            if version.max is not None and search_context.version > version.max:
+                continue
+            if typeshed_dir.name != "@python2" and version.in_python2:
                 continue
             if entry.is_dir():
                 seen = yield from _get_all_stub_files_from_directory(
@@ -223,35 +228,60 @@ def get_stub_file_name(
 
     # 5. typeshed
     versions = get_typeshed_versions(search_context.typeshed)
-    if (
-        top_level_name not in versions
-        or search_context.version < versions[top_level_name]
-    ):
+    if top_level_name not in versions:
+        return None
+    version = versions[top_level_name]
+    if search_context.version < version.min:
+        return None
+    if version.max is not None and search_context.version > version.max:
         return None
 
     if search_context.version[0] == 2:
         python2_dir = search_context.typeshed / "@python2"
         stub = _find_stub_in_dir(python2_dir, module_name)
-        if stub is not None:
+        if stub is not None or version.in_python2:
             return stub
 
     return _find_stub_in_dir(search_context.typeshed, module_name)
 
 
+class _VersionData(NamedTuple):
+    min: PythonVersion
+    max: Optional[PythonVersion]
+    # whether it is present in @python2
+    in_python2: bool
+
+
 @lru_cache()
-def get_typeshed_versions(typeshed: Path) -> Dict[str, PythonVersion]:
+def get_typeshed_versions(typeshed: Path) -> Dict[str, _VersionData]:
     versions = {}
+    python2_files = set(os.listdir(typeshed / "@python2"))
     with (typeshed / "VERSIONS").open() as f:
         for line in f:
             line = line.strip()
             if not line or line.startswith("#"):
                 continue
             module, version = line.split(": ")
-            major, minor = version.split(".")
-            versions[module] = (int(major), int(minor))
+            if "-" in version:
+                min_version_str, max_version_str = version.split("-")
+            else:
+                min_version_str = version
+                max_version_str = None
+            if max_version_str:
+                max_version = _parse_version(max_version_str)
+            else:
+                max_version = None
+            min_version = _parse_version(min_version_str)
+            python2_only = module in python2_files or module + ".pyi" in python2_files
+            versions[module] = _VersionData(min_version, max_version, python2_only)
     return versions
 
 
+def _parse_version(version: str) -> PythonVersion:
+    major, minor = version.split(".")
+    return (int(major), int(minor))
+
+
 def _find_stub_in_dir(stubdir: Path, module: ModulePath) -> Optional[Path]:
     if not module:
         init_name = stubdir / "__init__.pyi"

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,5 @@`
`1`	`1`	`from typing import overload`
	`2`	`+`
`2`	`3`	`@overload`
`3`	`4`	`def overloaded(x: int) -> None: ...`
`4`	`5`	`@overload`