Fix import expansion: keep full module paths, expand relative aliases

P1: parse_imports() no longer truncates to 2 components. Full module paths
like diff_diff.visualization._common are preserved. For 'from . import foo'
style relative imports, each alias is appended to the resolved base package
(e.g., diff_diff.visualization._event_study instead of diff_diff.visualization).

P2: Add 3 regression tests: submodule import truncation check,
relative import alias expansion, visualization __init__.py expansion
includes submodule files.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: igerber

.claude/scripts/openai_review.py

@@ -243,27 +243,26 @@ def parse_imports(file_path: str) -> "set[str]":
     for node in ast.walk(tree):
         if isinstance(node, ast.Import):
             for alias in node.names:
-                if alias.name.startswith("diff_diff"):
-                    # Extract the top-level module: diff_diff.linalg.foo -> diff_diff.linalg
-                    parts = alias.name.split(".")
-                    if len(parts) >= 2:
-                        imports.add(f"{parts[0]}.{parts[1]}")
+                if alias.name.startswith("diff_diff."):
+                    imports.add(alias.name)
         elif isinstance(node, ast.ImportFrom):
             if node.module and node.level == 0:
                 # Absolute import: from diff_diff.linalg import ...
-                if node.module.startswith("diff_diff"):
-                    parts = node.module.split(".")
-                    if len(parts) >= 2:
-                        imports.add(f"{parts[0]}.{parts[1]}")
+                if node.module.startswith("diff_diff."):
+                    imports.add(node.module)
             elif node.level > 0 and package:
-                # Relative import: from . import utils, from .linalg import solve_ols
+                # Relative import: from .foo import bar, or from . import foo
                 resolved = _resolve_relative_import(
                     package, node.module, node.level
                 )
                 if resolved and resolved.startswith("diff_diff"):
-                    parts = resolved.split(".")
-                    if len(parts) >= 2:
-                        imports.add(f"{parts[0]}.{parts[1]}")
+                    if node.module:
+                        # from .linalg import solve_ols → resolved = "diff_diff.linalg"
+                        imports.add(resolved)
+                    else:
+                        # from . import _event_study, _common → append each alias
+                        for alias in node.names:
+                            imports.add(f"{resolved}.{alias.name}")
     return imports
 
 

tests/test_openai_review.py

@@ -542,6 +542,32 @@ def test_ignores_non_diff_diff_imports(self, review_mod, tmp_path):
         imports = review_mod.parse_imports(str(test_file))
         assert imports == set()
 
+    def test_submodule_imports_not_truncated(self, review_mod, repo_root):
+        """Submodule imports should keep full path, not truncate to 2 components."""
+        path = os.path.join(repo_root, "diff_diff", "visualization", "_staggered.py")
+        if not os.path.isfile(path):
+            pytest.skip("diff_diff/visualization/_staggered.py not found")
+        imports = review_mod.parse_imports(path)
+        # Should include full submodule paths like diff_diff.visualization._common
+        has_submodule = any(
+            m.count(".") >= 2 for m in imports  # at least 3 components
+        )
+        assert has_submodule, (
+            f"Expected submodule imports (3+ components) but got: {imports}"
+        )
+
+    def test_relative_import_aliases_expanded(self, review_mod, repo_root):
+        """from . import _event_study should resolve to diff_diff.visualization._event_study."""
+        path = os.path.join(repo_root, "diff_diff", "visualization", "__init__.py")
+        if not os.path.isfile(path):
+            pytest.skip("diff_diff/visualization/__init__.py not found")
+        imports = review_mod.parse_imports(path)
+        # Should include individual submodule names, not just the package
+        submodules = [m for m in imports if m.startswith("diff_diff.visualization._")]
+        assert len(submodules) > 0, (
+            f"Expected visualization submodule imports but got: {imports}"
+        )
+
     def test_handles_syntax_error(self, review_mod, tmp_path, capsys):
         test_file = tmp_path / "bad.py"
         test_file.write_text("def foo(:\n  pass\n")
@@ -583,6 +609,18 @@ def test_deduplicates_against_changed_set(self, review_mod, repo_root):
         result = review_mod.expand_import_graph([bacon, linalg], repo_root)
         assert linalg not in [os.path.normpath(p) for p in result]
 
+    def test_visualization_init_includes_submodules(self, review_mod, repo_root):
+        """expand_import_graph on visualization/__init__.py should include submodules."""
+        path = os.path.join(repo_root, "diff_diff", "visualization", "__init__.py")
+        if not os.path.isfile(path):
+            pytest.skip("diff_diff/visualization/__init__.py not found")
+        result = review_mod.expand_import_graph([path], repo_root)
+        filenames = [os.path.basename(p) for p in result]
+        # Should include visualization submodules like _event_study.py, _staggered.py
+        assert any(f.startswith("_") and f.endswith(".py") for f in filenames), (
+            f"Expected visualization submodule files but got: {filenames}"
+        )
+
     def test_empty_input(self, review_mod, repo_root):
         assert review_mod.expand_import_graph([], repo_root) == []