Collect "from imports" as referencable names for docstrings

This makes it possible to use symbols made available via "from imports"
within the same module. This lessens the need to declare types manually.

Author: lagru

examples/example_pkg-stubs/__init__.pyi

@@ -10,3 +10,6 @@ __all__ = [
 
 class CustomException(Exception):
     pass
+
+class AnotherType:
+    pass

examples/example_pkg-stubs/_basic.pyi

@@ -11,6 +11,10 @@ from . import CustomException
 
 logger: Incomplete
 
+from configparser import ConfigParser
+
+from . import AnotherType
+
 __all__ = [
     "func_empty",
     "ExampleClass",
@@ -39,6 +43,7 @@ def func_use_from_elsewhere(
     a3: ExampleClass.NestedClass,
     a4: ExampleClass.NestedClass,
 ) -> tuple[CustomException, ExampleClass.NestedClass]: ...
+def func_use_from_import(a1: AnotherType, a2: ConfigParser) -> None: ...
 
 class ExampleClass:
 

examples/example_pkg/__init__.py

@@ -11,3 +11,7 @@
 
 class CustomException(Exception):
     pass
+
+
+class AnotherType:
+    pass

examples/example_pkg/_basic.py

@@ -88,6 +88,16 @@ def func_use_from_elsewhere(a1, a2, a3, a4):
     """
 
 
+def func_use_from_import(a1, a2):
+    """Check using symbols made available in this module with from imports.
+
+    Parameters
+    ----------
+    a1 : AnotherType
+    a2 : Cfg
+    """
+
+
 class ExampleClass:
     """Dummy.
 

src/docstub/_analysis.py

@@ -388,6 +388,18 @@ def visit_AnnAssign(self, node: cst.AnnAssign) -> bool:
             self._collect_type_annotation(stack)
         return False
 
+    def visit_ImportFrom(self, node: cst.ImportFrom) -> None:
+        """Collect "from import" targets as usable types within each module."""
+        for import_alias in node.names:
+            if cstm.matches(import_alias, cstm.ImportStar()):
+                continue
+            name = import_alias.evaluated_alias
+            if name is None:
+                name = import_alias.evaluated_name
+            assert isinstance(name, str)
+            stack = [*self._stack, name]
+            self._collect_type_annotation(stack)
+
     def _collect_type_annotation(self, stack):
         """Collect an importable type annotation.
 

tests/test_analysis.py

@@ -49,7 +49,6 @@ def _module_factory(src, module_name):
 
 
 class Test_TypeCollector:
-
     def test_classes(self, module_factory):
         module_path = module_factory(
             src=dedent(
@@ -100,6 +99,41 @@ def test_ignores_assigns(self, module_factory, src):
         imports = TypeCollector.collect(file=module_path)
         assert len(imports) == 0
 
+    @pytest.mark.parametrize(
+        "src",
+        [
+            "from calendar import Aug",
+            "from . import Aug",
+            "from calendar import August as Aug",
+            "from . import Agust as Aug",
+        ],
+    )
+    def test_from_import(self, module_factory, src):
+        module_path = module_factory(src=src, module_name="sub.module")
+        imports = TypeCollector.collect(file=module_path)
+        assert len(imports) == 1
+        assert imports == {
+            "sub.module.Aug": KnownImport(import_path="sub.module", import_name="Aug")
+        }
+
+    @pytest.mark.parametrize(
+        "src",
+        [
+            "from calendar import Aug, Dec",
+            "from . import Aug, Dec",
+            "from calendar import August as Aug, December as Dec",
+            "from . import August as Aug, December as Dec",
+        ],
+    )
+    def test_from_import_multiple(self, module_factory, src):
+        module_path = module_factory(src=src, module_name="sub.module")
+        imports = TypeCollector.collect(file=module_path)
+        assert len(imports) == 2
+        assert imports == {
+            "sub.module.Aug": KnownImport(import_path="sub.module", import_name="Aug"),
+            "sub.module.Dec": KnownImport(import_path="sub.module", import_name="Dec"),
+        }
+
 
 class Test_TypeMatcher:
     type_prefixes = {  # noqa: RUF012