get_all_stub_files should also return the module names

JelleZijlstra · JelleZijlstra · commit 2012973ba2df · 2017-05-22T16:46:54.000-07:00
diff --git a/tests/test.py b/tests/test.py
@@ -45,9 +45,9 @@ def test_subdir(self) -> None:
     def test_get_all_stub_files(self) -> None:
         all_stubs = typeshed_client.get_all_stub_files(version=(2, 7), typeshed_dir=TEST_TYPESHED)
         self.assertEqual(set(all_stubs), {
-            TEST_TYPESHED / 'stdlib/2/lib.pyi',
-            TEST_TYPESHED / 'stdlib/2and3/conditions.pyi',
-            TEST_TYPESHED / 'stdlib/2and3/shared.pyi',
+            ('lib', TEST_TYPESHED / 'stdlib/2/lib.pyi'),
+            ('conditions', TEST_TYPESHED / 'stdlib/2and3/conditions.pyi'),
+            ('shared', TEST_TYPESHED / 'stdlib/2and3/shared.pyi'),
         })
 
 
diff --git a/typeshed_client/finder.py b/typeshed_client/finder.py
@@ -72,8 +72,12 @@ def parse_stub_file(path: Path) -> ast3.AST:
 
 
 def get_all_stub_files(version: Tuple[int, int] = sys.version_info[:2],
-                       typeshed_dir: Optional[Path] = None) -> Iterable[Path]:
-    """Returns paths to all stub files for a given Python version."""
+                       typeshed_dir: Optional[Path] = None) -> Iterable[Tuple[str, Path]]:
+    """Returns paths to all stub files for a given Python version.
+
+    Returns pairs of (module name, module path).
+
+    """
     if typeshed_dir is None:
         typeshed_dir = find_typeshed()
     seen: Set[str] = set()
@@ -86,5 +90,15 @@ def get_all_stub_files(version: Tuple[int, int] = sys.version_info[:2],
                     continue
                 if relative_path in seen:
                     continue
-                yield full_path
+                yield _path_to_module(relative_path), full_path
                 seen.add(relative_path)
+
+
+def _path_to_module(path: Path) -> str:
+    """Returns the module name corresponding to a file path."""
+    parts = path.parts
+    if parts[-1] == '__init__.pyi':
+        parts = parts[-1]
+    if parts[-1].endswith('.pyi'):
+        parts = parts[:-1] + (parts[-1].rstrip('.pyi'),)
+    return '.'.join(parts)
