Generalize AST import detection through shared symbol helper

Co-authored-by: Shri Sukhani <shrisukhani@users.noreply.github.com>

Author: cursoragent

tests/ast_import_utils.py

@@ -1,25 +1,29 @@
 import ast
 
 
-def imports_collect_function_sources(module_text: str) -> bool:
+def imports_symbol_from_module(module_text: str, module: str, symbol: str) -> bool:
     module_ast = ast.parse(module_text)
     for node in module_ast.body:
         if not isinstance(node, ast.ImportFrom):
             continue
-        if node.module != "tests.ast_function_source_utils":
+        if node.module != module:
             continue
-        if any(alias.name == "collect_function_sources" for alias in node.names):
+        if any(alias.name == symbol for alias in node.names):
             return True
     return False
 
 
+def imports_collect_function_sources(module_text: str) -> bool:
+    return imports_symbol_from_module(
+        module_text,
+        module="tests.ast_function_source_utils",
+        symbol="collect_function_sources",
+    )
+
+
 def imports_imports_collect_function_sources(module_text: str) -> bool:
-    module_ast = ast.parse(module_text)
-    for node in module_ast.body:
-        if not isinstance(node, ast.ImportFrom):
-            continue
-        if node.module != "tests.ast_import_utils":
-            continue
-        if any(alias.name == "imports_collect_function_sources" for alias in node.names):
-            return True
-    return False
+    return imports_symbol_from_module(
+        module_text,
+        module="tests.ast_import_utils",
+        symbol="imports_collect_function_sources",
+    )

tests/test_ast_import_utils.py

@@ -1,6 +1,7 @@
 import pytest
 
 from tests.ast_import_utils import (
+    imports_symbol_from_module,
     imports_collect_function_sources,
     imports_imports_collect_function_sources,
 )
@@ -26,6 +27,38 @@ def test_imports_collect_function_sources_ignores_non_matching_imports():
     assert imports_collect_function_sources(module_text) is False
 
 
+def test_imports_symbol_from_module_detects_expected_symbol():
+    module_text = (
+        "from tests.ast_import_utils import imports_collect_function_sources\n"
+        "imports_collect_function_sources('dummy')\n"
+    )
+
+    assert (
+        imports_symbol_from_module(
+            module_text,
+            module="tests.ast_import_utils",
+            symbol="imports_collect_function_sources",
+        )
+        is True
+    )
+
+
+def test_imports_symbol_from_module_ignores_unrelated_symbols():
+    module_text = (
+        "from tests.ast_import_utils import imports_collect_function_sources\n"
+        "from tests.ast_import_utils import imports_imports_collect_function_sources\n"
+    )
+
+    assert (
+        imports_symbol_from_module(
+            module_text,
+            module="tests.ast_import_utils",
+            symbol="missing_symbol",
+        )
+        is False
+    )
+
+
 def test_imports_collect_function_sources_supports_aliased_import():
     module_text = (
         "from tests.ast_function_source_utils import collect_function_sources as cfs\n"