Repository: symbol-management/symbol-exporter

Issue:
Add top level imports to volume of package/module
Top level imports must exist to use the package or module so we should add them to the volume also.

Side note: Once we handle imports in functions we will need to add them to the volume of that function as well.


Discussion:

Patch:
diff --git a/symbol_exporter/ast_symbol_extractor.py b/symbol_exporter/ast_symbol_extractor.py
index 39caf11..39d1c18 100644
--- a/symbol_exporter/ast_symbol_extractor.py
+++ b/symbol_exporter/ast_symbol_extractor.py
@@ -97,6 +97,7 @@ class SymbolFinder(ast.NodeVisitor):
     def visit_Import(self, node: ast.Import) -> Any:
         self.imported_symbols.extend(k.name for k in node.names)
         for k in node.names:
+            self._add_symbol_to_volume(self._symbol_stack_to_symbol_name(), k.name, lineno=node.lineno)
             if not k.asname:
                 self._add_symbol_to_surface_area(SymbolType.IMPORT, symbol=k.name, shadows=k.name)
         self.generic_visit(node)
@@ -108,9 +109,11 @@ class SymbolFinder(ast.NodeVisitor):
                 module_name = f"{node.module}.{k.name}" if node.module else k.name
                 self.aliases[k.name] = module_name
                 self.imported_symbols.append(module_name)
+                # TODO: clean up if statements with common module_name var between forks
                 if not k.asname:
                     if not relative_import:
                         self._add_symbol_to_surface_area(SymbolType.IMPORT, symbol=k.name, shadows=module_name)
+                    # TODO: handle import in volume for relative imports
                     else:
                         self._add_symbol_to_surface_area(
                             SymbolType.RELATIVE_IMPORT,
@@ -119,15 +122,20 @@ class SymbolFinder(ast.NodeVisitor):
                             level=node.level,
                             module=self._module_name,
                         )
+                if not relative_import:
+                    self._add_symbol_to_volume(self._symbol_stack_to_symbol_name(), module_name, lineno=node.lineno)
             else:
                 if not relative_import:
                     self._add_symbol_to_star_imports(node.module, symbol_type=SymbolType.STAR_IMPORT)
+                # TODO: handle import in volume for relative imports
                 else:
                     self._add_symbol_to_relative_star_imports(
                         node.module,
                         symbol_type=SymbolType.RELATIVE_STAR_IMPORT,
                         level=node.level,
                     )
+                if not relative_import:
+                    self._add_symbol_to_volume(self._symbol_stack_to_symbol_name(), node.module, lineno=node.lineno)
         self._relative_import_stack.append(node.level)
         self.generic_visit(node)
         self._relative_import_stack.pop(-1)
@@ -266,11 +274,11 @@ class SymbolFinder(ast.NodeVisitor):
             return None
 
     def _add_symbol_to_surface_area(self, symbol_type: SymbolType, symbol, **kwargs):
-        full_symbol_name = (
-            f"{self._module_name}.{symbol}"
-            if symbol_type in (SymbolType.IMPORT, SymbolType.RELATIVE_IMPORT)
-            else symbol
-        )
+        imports = (SymbolType.IMPORT, SymbolType.RELATIVE_IMPORT)
+        # imports inside namespaces other than top level aren't on the surface area
+        if symbol_type in imports and len(self.current_symbol_stack) > 1:
+            return
+        full_symbol_name = f"{self._module_name}.{symbol}" if symbol_type in imports else symbol
         self._symbols[full_symbol_name] = dict(type=symbol_type, data=kwargs)
 
     def _add_symbol_to_volume(self, surface_symbol, volume_symbol, lineno):


Test patch:
diff --git a/tests/test_ast_package_symbol_extractor.py b/tests/test_ast_package_symbol_extractor.py
index 555dd47..66ab18e 100644
--- a/tests/test_ast_package_symbol_extractor.py
+++ b/tests/test_ast_package_symbol_extractor.py
@@ -4,9 +4,9 @@ from symbol_exporter.ast_package_symbol_extractor import DirectorySymbolFinder
 from symbol_exporter.ast_symbol_extractor import SymbolType
 
 expected = {
-    "numpy.version": {"type": SymbolType.MODULE, "data": {}},
+    "numpy.version": {"type": SymbolType.MODULE, "data": {"symbols_in_volume": {"os": {"line number": [1]}}}},
     "numpy.version.get_versions": {"type": SymbolType.FUNCTION, "data": {"lineno": 5}},
-    "numpy": {"type": SymbolType.PACKAGE, "data": {}},
+    "numpy": {"type": SymbolType.PACKAGE, "data": {"symbols_in_volume": {"requests": {"line number": [5]}}}},
     "numpy.numeric": {
         "type": SymbolType.RELATIVE_IMPORT,
         "data": {"shadows": "numpy.core.numeric"},
@@ -47,7 +47,7 @@ expected = {
         "type": SymbolType.RELATIVE_STAR_IMPORT,
         "data": {"imports": ["numpy.core"]},
     },
-    "numpy.core": {"type": SymbolType.PACKAGE, "data": {}},
+    "numpy.core": {"type": SymbolType.PACKAGE, "data": {"symbols_in_volume": {"json": {"line number": [8]}}}},
     "numpy.core.~~RELATIVE~~.*": {
         "type": SymbolType.RELATIVE_STAR_IMPORT,
         "data": {"imports": ["numpy.core.numeric"]},
diff --git a/tests/test_ast_symbol_extractor.py b/tests/test_ast_symbol_extractor.py
index f991d56..1902cd0 100644
--- a/tests/test_ast_symbol_extractor.py
+++ b/tests/test_ast_symbol_extractor.py
@@ -18,8 +18,8 @@ def test_package_name():
     z = process_code_str(code, module_name="__init__")
     assert z.symbols == {
         "__init__": {
+            "data": {"symbols_in_volume": {"abc.xyz": {"line number": [2]}}},
             "type": SymbolType.PACKAGE,
-            "data": {},
         },
         "__init__.xyz": {
             "type": SymbolType.IMPORT,
@@ -35,8 +35,8 @@ def test_fully_qualified_package_names():
     z = process_code_str(code, module_name="mm.__init__")
     assert z.symbols == {
         "mm.__init__": {
+            "data": {"symbols_in_volume": {"abc.xyz": {"line number": [2]}}},
             "type": SymbolType.PACKAGE,
-            "data": {},
         },
         "mm.__init__.xyz": {
             "type": SymbolType.IMPORT,
@@ -52,8 +52,8 @@ def test_fully_qualified_module_names():
     z = process_code_str(code, module_name="mm.core")
     assert z.symbols == {
         "mm.core": {
+            "data": {"symbols_in_volume": {"abc.xyz": {"line number": [2]}}},
             "type": SymbolType.MODULE,
-            "data": {},
         },
         "mm.core.xyz": {
             "type": SymbolType.IMPORT,
@@ -75,8 +75,8 @@ def test_from_import_attr_access():
     assert z.used_symbols == {"abc.xyz.i"}
     assert z.symbols == {
         "mm": {
+            "data": {"symbols_in_volume": {"abc.xyz": {"line number": [2]}}},
             "type": SymbolType.MODULE,
-            "data": {},
         },
         "mm.xyz": {
             "type": SymbolType.IMPORT,
@@ -92,6 +92,28 @@ def test_from_import_attr_access():
     }
 
 
+def test_import_in_function():
+    code = """
+    def f():
+        from abc import xyz
+        return xyz.i
+    """
+    z = process_code_str(code)
+    assert z.symbols == {
+        "mm": {
+            "type": SymbolType.MODULE,
+            "data": {},
+        },
+        "mm.f": {
+            "type": SymbolType.FUNCTION,
+            "data": {
+                "lineno": 2,
+                "symbols_in_volume": {"abc.xyz.i": {"line number": [4]}, "abc.xyz": {"line number": [3]}},
+            },
+        },
+    }
+
+
 def test_alias_import():
     code = """
     from abc import xyz as l
@@ -106,7 +128,7 @@ def test_alias_import():
     assert z.symbols == {
         "mm": {
             "type": SymbolType.MODULE,
-            "data": {},
+            "data": {"symbols_in_volume": {"abc.xyz": {"line number": [2]}}},
         },
         "mm.l": {
             "type": SymbolType.IMPORT,
@@ -137,7 +159,7 @@ def test_import_with_and_without_alias_exposes_import_and_alias():
     assert z.symbols == {
         "mm": {
             "type": SymbolType.MODULE,
-            "data": {},
+            "data": {"symbols_in_volume": {"abc.xyz": {"line number": [2, 3]}}},
         },
         "mm.f": {
             "type": SymbolType.FUNCTION,
@@ -171,7 +193,7 @@ def test_calls():
     assert z.symbols == {
         "mm": {
             "type": SymbolType.MODULE,
-            "data": {},
+            "data": {"symbols_in_volume": {"numpy": {"line number": [2]}}},
         },
         "mm.f": {
             "type": SymbolType.FUNCTION,
@@ -203,7 +225,7 @@ def test_constant():
     assert z.symbols == {
         "mm": {
             "type": SymbolType.MODULE,
-            "data": {"symbols_in_volume": {"numpy.ones": {"line number": [4]}}},
+            "data": {"symbols_in_volume": {"numpy.ones": {"line number": [4]}, "numpy": {"line number": [2]}}},
         },
         "mm.np": {
             "type": SymbolType.IMPORT,
@@ -230,7 +252,7 @@ def test_class():
     assert z.symbols == {
         "mm": {
             "type": SymbolType.MODULE,
-            "data": {},
+            "data": {"symbols_in_volume": {"numpy": {"line number": [2]}}},
         },
         "mm.np": {
             "type": SymbolType.IMPORT,
@@ -263,7 +285,7 @@ def test_class_method():
     assert z.symbols == {
         "mm": {
             "type": SymbolType.MODULE,
-            "data": {},
+            "data": {"symbols_in_volume": {"numpy": {"line number": [2]}}},
         },
         "mm.ABC": {
             "type": SymbolType.CLASS,
@@ -314,8 +336,16 @@ def test_import_adds_symbols():
             "data": {"shadows": "ghi"},
         },
         "mm": {
+            "data": {
+                "symbols_in_volume": {
+                    "abc.xyz": {"line number": [3]},
+                    "ggg.efg": {"line number": [4]},
+                    "mm.ghi": {"line number": [5]},
+                    "numpy": {"line number": [2]},
+                    "numpy.ones": {"line number": [7]},
+                }
+            },
             "type": SymbolType.MODULE,
-            "data": {"symbols_in_volume": {"numpy.ones": {"line number": [7]}}},
         },
         "mm.z": {
             "type": SymbolType.CONSTANT,
@@ -344,7 +374,7 @@ def test_star_import():
         },
         "mm": {
             "type": SymbolType.MODULE,
-            "data": {},
+            "data": {"symbols_in_volume": {"abc": {"line number": [3]}, "numpy": {"line number": [2]}}},
         },
     }
 
@@ -453,6 +483,9 @@ def test_undeclared_symbols():
             "type": SymbolType.MODULE,
             "data": {
                 "symbols_in_volume": {
+                    "abc": {"line number": [4]},
+                    "numpy": {"line number": [2]},
+                    "xyz": {"line number": [5]},
                     "numpy.ones": {"line number": [8]},
                     "twos": {"line number": [9]},
                 },
@@ -486,7 +519,7 @@ def test_imported_symbols_not_treated_as_undeclared():
     assert z.symbols == {
         "mm": {
             "type": SymbolType.MODULE,
-            "data": {"symbols_in_volume": {"abc.twos": {"line number": [4]}}},
+            "data": {"symbols_in_volume": {"abc.twos": {"line number": [2, 4]}}},
         },
         "mm.twos": {
             "type": SymbolType.IMPORT,
@@ -514,7 +547,9 @@ def test_builtin_symbols_not_treated_as_undeclared():
     assert z.symbols == {
         "mm": {
             "type": SymbolType.MODULE,
-            "data": {"symbols_in_volume": {"len": {"line number": [4]}}},
+            "data": {
+                "symbols_in_volume": {"len": {"line number": [4]}, "abc.twos": {"line number": [2]}},
+            },
         },
         "mm.twos": {
             "type": SymbolType.IMPORT,
@@ -545,7 +580,9 @@ def test_functions_not_treated_as_undeclared():
     assert z.symbols == {
         "mm": {
             "type": SymbolType.MODULE,
-            "data": {"symbols_in_volume": {"mm.f": {"line number": [7]}}},
+            "data": {
+                "symbols_in_volume": {"mm.f": {"line number": [7]}, "abc.twos": {"line number": [2]}},
+            },
         },
         "mm.f": {
             "type": SymbolType.FUNCTION,
@@ -575,6 +612,7 @@ def test_attr_assignment():
         "mm": {
             "data": {
                 "symbols_in_volume": {
+                    "abc.twos": {"line number": [2]},
                     "abc.twos.three": {"line number": [4]},
                     "abc.twos.four": {"line number": [5]},
                 }
@@ -656,6 +694,7 @@ def test_symbols_in_volume_names():
             "data": {
                 "symbols_in_volume": {
                     "ast.Try": {"line number": [4]},
+                    "mm.ast": {"line number": [2]},
                     "mm.z": {"line number": [5]},
                 }
             },