add ImportFrom support DefinitionFinder and tests

iyume · iyume · commit 3cf3067d61b5 · 2022-10-29T15:03:20.000+08:00
diff --git a/.vscode/settings.json b/.vscode/settings.json
@@ -1,11 +1,11 @@
 {
     "python.defaultInterpreterPath": ".venv/bin/python",
     "python.formatting.provider": "black",
-    "[python]": {
-        "editor.codeActionsOnSave": {
-            "source.organizeImports": true
-        }
-    },
+    // "[python]": {
+    //     "editor.codeActionsOnSave": {
+    //         "source.organizeImports": true
+    //     }
+    // },
     "python.analysis.typeCheckingMode": "basic",
     "python.testing.pytestArgs": [
         "tests"
diff --git a/nb_autodoc/analyzers/definitionfinder.py b/nb_autodoc/analyzers/definitionfinder.py
@@ -38,11 +38,11 @@ class ClassDefData:
 
 
 @dataclass
-class External:
+class ImportFromData:
     order: int
-    varname: str
+    name: str  # import asname
     module: str
-    name: str
+    orig_name: str
 
 
 def signature_from_ast(node: ast.FunctionDef) -> Signature:
@@ -67,17 +67,19 @@ class DefinitionFinder:
 
     The following patterns bind names:
 
-        - class or function definition
-        - assignment expression (:=)
-        - targets that create new variable
-        - import statement
-        - compound statement body (if, while, for, try, with, match)
-        - See: https://docs.python.org/3/reference/executionmodel.html#naming-and-binding
+    - class or function definition
+    - assignment expression (:=)
+    - targets that create new variable
+    - import statement
+    - compound statement body (if, while, for, try, with, match)
+    - See: https://docs.python.org/3/reference/executionmodel.html#naming-and-binding
+
+    We bind names on `from...import`, `new variable assignment`, `function or class`
 
     A name can be classified into:
 
-        - definition: assign, function and class
-        - external: ImportFrom matches user module
+    - definition: assign, function and class
+    - external: ImportFrom matches user module
 
     **Variable comment:**
 
@@ -101,11 +103,11 @@ class DefinitionFinder:
     will be assigned to each name.
     """
 
-    def __init__(self, name: str, package: Optional[str]) -> None:
+    def __init__(self, *, package: Optional[str]) -> None:
         # arg should all be optional (as analysis feature)
-        self.name = name
-        """Module name. Determind external import."""
-        self.package = package
+        # self.name = name
+        # """Module name. Determind external import."""
+        self.package = package  # maybe null
         """Package name. Resolve relative import."""
         self.next_stmt: Optional[ast.stmt] = None
         self.current_classes: List[ClassDefData] = []
@@ -207,10 +209,15 @@ def visit_AnnAssign(self, node: ast.AnnAssign) -> None:
         return self.visit_Assign(node)  # type: ignore
 
     def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
+        if self.current_function:
+            return
         absoluate_module = resolve_name(node, self.package)
-        if absoluate_module.split(".", 1)[0] == self.name.split(".", 1):
-            for alias in node.names:
-                varname = alias.asname or alias.name
+        scope = self.get_current_scope()
+        for alias in node.names:
+            varname = alias.asname or alias.name
+            scope[varname] = ImportFromData(
+                next(self.counter), varname, absoluate_module, alias.name
+            )
 
     # def visit_Expr(self, node: ast.Expr) -> None:
     #     # bound docstring on previous assign
diff --git a/nb_autodoc/analyzers/utils.py b/nb_autodoc/analyzers/utils.py
@@ -121,6 +121,7 @@ def get_assign_names(
 def resolve_name(
     name_or_import: t.Union[ast.ImportFrom, str], package: t.Optional[str] = None
 ) -> str:
+    """Resolve a relative module name to an absolute one."""
     if isinstance(name_or_import, ast.ImportFrom):
         name_or_import = "." * name_or_import.level + (name_or_import.module or "")
     return imp_resolve_name(name_or_import, package)
diff --git a/pyproject.toml b/pyproject.toml
@@ -30,6 +30,7 @@ target-version = ["py37"]
 profile = "black"
 extra_standard_library = ["typing_extensions"]
 skip_gitignore = true
+skip_glob = ["tests/analyzerdata/*"]
 
 [tool.pytest.ini_options]
 log_cli = true  # for temporary debug
diff --git a/tests/analyzerdata/simple-definition-ast.py b/tests/analyzerdata/simple-definition-ast.py
@@ -2,10 +2,10 @@
 # fmt: off
 
 import os
-from pathlib import Path
+from pathlib import Path as Path_rename
 
-# give module and package name for DefinitionFinder to analyze following import stmt
-from mypkg import ext_A, ext_fa
+# give package name to DefinitionFinder to analyze following import stmt
+from mypkg import ext_A as ext_A_rename, ext_fa
 from mypkg.pkg import ext_B, ext_fb
 
 from .util import ext_fc
diff --git a/tests/test_analyzers/test_definitionfinder.py b/tests/test_analyzers/test_definitionfinder.py
@@ -6,108 +6,127 @@
     ClassDefData,
     DefinitionFinder,
     FunctionDefData,
+    ImportFromData,
 )
 from nb_autodoc.analyzers.utils import ast_parse
 
 
 def get_analyzer_data(filename: str) -> str:
-    return open(Path(__file__).parent / "analyzerdata" / filename).read()
+    return open(Path(__file__).parent.parent / "analyzerdata" / filename).read()
 
 
 class TestDefinitionFinder:
     def test_simple_data(self):
         code = get_analyzer_data("simple-definition-ast.py")
         module = ast_parse(code)
-        visitor = DefinitionFinder("mypkg.pkg", "mypkg.pkg.pkg")
+        visitor = DefinitionFinder(package="mypkg.pkg.pkg")
         visitor.visit(module)
         del code, module
         module_freevars = visitor.freevars
         # duplicated from repr(module_freevars)
         # notice that AssignData.annotation is uncomparable `ast.Expression`
         # so until unparser implement (maybe py3.8), annotation test is unavailable
         assert module_freevars == {
+            "Path_rename": ImportFromData(
+                order=0, name="Path_rename", module="pathlib", orig_name="Path"
+            ),
+            "ext_A_rename": ImportFromData(
+                order=1, name="ext_A_rename", module="mypkg", orig_name="ext_A"
+            ),
+            "ext_fa": ImportFromData(
+                order=2, name="ext_fa", module="mypkg", orig_name="ext_fa"
+            ),
+            "ext_B": ImportFromData(
+                order=3, name="ext_B", module="mypkg.pkg", orig_name="ext_B"
+            ),
+            "ext_fb": ImportFromData(
+                order=4, name="ext_fb", module="mypkg.pkg", orig_name="ext_fb"
+            ),
+            "ext_fc": ImportFromData(
+                order=5, name="ext_fc", module="mypkg.pkg.pkg.util", orig_name="ext_fc"
+            ),
             "a": AssignData(
-                order=0, name="a", type_comment="int", docstring="a docstring"
+                order=6, name="a", type_comment="int", docstring="a docstring"
             ),
-            "a2": AssignData(order=1, name="a2", type_comment="int", docstring=None),
-            "a3": AssignData(order=2, name="a3", type_comment=None, docstring=None),
+            "a2": AssignData(order=7, name="a2", type_comment="int", docstring=None),
+            "a3": AssignData(order=8, name="a3", type_comment=None, docstring=None),
             "b": AssignData(
-                order=3, name="b", type_comment="int", docstring="b docstring"
+                order=9, name="b", type_comment="int", docstring="b docstring"
             ),
-            "fa": FunctionDefData(order=4, name="fa"),
+            "fa": FunctionDefData(order=10, name="fa"),
             "c": AssignData(
-                order=5, name="c", type_comment=None, docstring="c and d docstring"
+                order=11, name="c", type_comment=None, docstring="c and d docstring"
             ),
             "d": AssignData(
-                order=6, name="d", type_comment=None, docstring="c and d docstring"
+                order=12, name="d", type_comment=None, docstring="c and d docstring"
             ),
             "a1": AssignData(
-                order=7, name="a1", type_comment=None, docstring="abcde11111 docstring"
+                order=13, name="a1", type_comment=None, docstring="abcde11111 docstring"
             ),
             "b1": AssignData(
-                order=8, name="b1", type_comment=None, docstring="abcde11111 docstring"
+                order=14, name="b1", type_comment=None, docstring="abcde11111 docstring"
             ),
             "c1": AssignData(
-                order=9, name="c1", type_comment=None, docstring="abcde11111 docstring"
+                order=15, name="c1", type_comment=None, docstring="abcde11111 docstring"
             ),
             "d1": AssignData(
-                order=10, name="d1", type_comment=None, docstring="abcde11111 docstring"
+                order=16, name="d1", type_comment=None, docstring="abcde11111 docstring"
             ),
             "e1": AssignData(
-                order=11, name="e1", type_comment=None, docstring="abcde11111 docstring"
+                order=17, name="e1", type_comment=None, docstring="abcde11111 docstring"
             ),
             "A": ClassDefData(
-                order=12,
+                order=18,
                 name="A",
                 freevars={
                     "a": AssignData(
-                        order=13, name="a", type_comment=None, docstring=None
+                        order=19, name="a", type_comment=None, docstring=None
                     )
                 },
                 instance_vars={},
                 methods={},
             ),
             "B": ClassDefData(
-                order=14,
+                order=20,
                 name="B",
                 freevars={
                     "a": AssignData(
-                        order=15, name="a", type_comment=None, docstring="B.a docstring"
+                        order=21, name="a", type_comment=None, docstring="B.a docstring"
                     )
                 },
                 instance_vars={},
                 methods={},
             ),
             "B1": ClassDefData(
-                order=16,
+                order=22,
                 name="B1",
                 freevars={
                     "a": AssignData(
-                        order=17, name="a", type_comment=None, docstring=None
+                        order=23, name="a", type_comment=None, docstring=None
                     ),
-                    "__init__": FunctionDefData(order=18, name="__init__"),
-                    "b": FunctionDefData(order=19, name="b"),
+                    "__init__": FunctionDefData(order=24, name="__init__"),
+                    "b": FunctionDefData(order=25, name="b"),
                 },
                 instance_vars={},
                 methods={},
             ),
             "C": ClassDefData(
-                order=20,
+                order=26,
                 name="C",
                 freevars={
                     "a": AssignData(
-                        order=21,
+                        order=27,
                         name="a",
                         type_comment=None,
                         docstring="C.a classvar docstring",
                     ),
-                    "__init__": FunctionDefData(order=22, name="__init__"),
+                    "__init__": FunctionDefData(order=28, name="__init__"),
                     "_A": ClassDefData(
-                        order=27,
+                        order=33,
                         name="_A",
                         freevars={
                             "_a": AssignData(
-                                order=28,
+                                order=34,
                                 name="_a",
                                 type_comment=None,
                                 docstring="nested OK",
@@ -119,22 +138,22 @@ def test_simple_data(self):
                 },
                 instance_vars={
                     "a": AssignData(
-                        order=23,
+                        order=29,
                         name="a",
                         type_comment="int | None",
                         docstring="C instance var a/b docstring",
                     ),
                     "b": AssignData(
-                        order=24,
+                        order=30,
                         name="b",
                         type_comment="int | None",
                         docstring="C instance var a/b docstring",
                     ),
                     "c": AssignData(
-                        order=25, name="c", type_comment=None, docstring=None
+                        order=31, name="c", type_comment=None, docstring=None
                     ),
                     "d": AssignData(
-                        order=26,
+                        order=32,
                         name="d",
                         type_comment=None,
                         docstring="C instance var d docstring",
@@ -147,7 +166,7 @@ def test_simple_data(self):
     def test_assigndata_override(self):
         code = get_analyzer_data("assigndata-override-ast.py")
         module = ast_parse(code)
-        visitor = DefinitionFinder("<test>", "<test>")
+        visitor = DefinitionFinder(package="<test>")
         visitor.visit(module)
         module_freevars = visitor.freevars
         assert module_freevars == {
diff --git a/tests/test_pkg/typing.py b/tests/test_pkg/typing.py
@@ -1,6 +1,5 @@
 from typing import Dict, Optional
 
-
 T_foo = Dict[str, str]
 """typing variable T_foo
 类型版本: 1.1.0+
