Support resolving via "from ... import ..." in normal python modules

typeshed_client/finder.py

@@ -6,6 +6,7 @@
 import subprocess
 import sys
 from functools import lru_cache
+from importlib.util import find_spec
 from pathlib import Path
 from typing import (
     TYPE_CHECKING,
@@ -103,6 +104,12 @@ def get_stub_file(
     return get_stub_file_name(ModulePath(tuple(module_name.split("."))), search_context)
 
 
+def get_py_module_file(module_name: ModulePath) -> Optional[Path]:
+    """Return the path to the python file for this module, if any."""
+    spec = find_spec(".".join(module_name))
+    return Path(spec.origin) if spec and spec.origin else None
+
+
 def get_stub_ast(
     module_name: str, *, search_context: Optional[SearchContext] = None
 ) -> Optional[ast.Module]:

typeshed_client/parser.py

@@ -12,6 +12,7 @@
     NamedTuple,
     NoReturn,
     Optional,
+    Set,
     Tuple,
     Type,
     Union,
@@ -438,3 +439,56 @@ def _warn(message: str, ctx: SearchContext) -> None:
         raise InvalidStub(message)
     else:
         log.warning(message)
+
+
+def get_py_imported_name_sources(
+    module_name: ModulePath,
+) -> Tuple[Dict[str, str], Set[str]]:
+    """Given a python module path, return named and star import sources."""
+    path = finder.get_py_module_file(module_name)
+    if path is None:
+        return {}, set()
+    ast = parse_stub_file(path)
+    visitor = _PyImportedSourcesExtractor(module_name)
+    visitor.visit(ast)
+    return visitor.named_imports, visitor.star_imports
+
+
+class _PyImportedSourcesExtractor(ast.NodeVisitor):
+    """Extract imported sources from a normal python module."""
+
+    def __init__(self, module_path: ModulePath) -> None:
+        self.module_path = module_path
+        self.named_imports: Dict[str, str] = {}
+        self.star_imports: Set[str] = set()
+
+    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
+        if node.level:
+            # Relative import
+            module_path = self.module_path[: len(self.module_path) - node.level + 1]
+            if node.module:
+                module_path = module_path + tuple(node.module.split("."))
+            import_module_path: str = ".".join(module_path)
+        else:
+            # Absolute import
+            import_module_path = node.module  # type: ignore[assignment]  # always str when level=0
+
+        if len(node.names) == 1 and node.names[0].name == "*":
+            # Star import
+            self.star_imports.add(import_module_path)
+        else:
+            # Named import
+            for alias in node.names:
+                self.named_imports[alias.asname or alias.name] = import_module_path
+
+    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
+        # Skip visiting function definitions
+        pass
+
+    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
+        # Skip visiting async function definitions
+        pass
+
+    def visit_ClassDef(self, node: ast.ClassDef) -> None:
+        # Skip visiting class definitions
+        pass

typeshed_client/resolver.py

@@ -15,10 +15,15 @@ class ImportedInfo(NamedTuple):
 
 
 class Resolver:
-    def __init__(self, search_context: Optional[SearchContext] = None) -> None:
+    def __init__(
+        self,
+        search_context: Optional[SearchContext] = None,
+        search_py_files: bool = False,
+    ) -> None:
         if search_context is None:
             search_context = get_search_context()
         self.ctx = search_context
+        self.search_py_files = search_py_files
         self._module_cache: Dict[ModulePath, Module] = {}
 
     def get_module(self, module_name: ModulePath) -> "Module":
@@ -34,7 +39,32 @@ def get_module(self, module_name: ModulePath) -> "Module":
 
     def get_name(self, module_name: ModulePath, name: str) -> ResolvedName:
         module = self.get_module(module_name)
-        return module.get_name(name, self)
+        resolved_name = module.get_name(name, self)
+        if resolved_name is None and self.search_py_files:
+            resolved_name = self._get_py_imported_name(module_name, name)
+        return resolved_name
+
+    def _get_py_imported_name(self, module_name: ModulePath, name: str) -> ResolvedName:
+        resolved_name = None
+        try:
+            named_imports, star_imports = parser.get_py_imported_name_sources(
+                module_name
+            )
+            for star_import in star_imports:
+                import_source = ModulePath(tuple(star_import.split(".")))
+                if import_source != module_name:
+                    # Try to resolve stub from star import
+                    resolved_name = self.get_name(import_source, name)
+                if resolved_name is not None:
+                    break
+            if resolved_name is None and name in named_imports:
+                import_source = ModulePath(tuple(named_imports[name].split(".")))
+                if import_source != module_name:
+                    # Try to resolve stub from named import
+                    resolved_name = self.get_name(import_source, name)
+        except Exception:
+            pass
+        return resolved_name
 
     def get_fully_qualified_name(self, name: str) -> ResolvedName:
         """Public API."""