diff --git a/code2flow/model.py b/code2flow/model.py index adb4a36..525adae 100644 --- a/code2flow/model.py +++ b/code2flow/model.py @@ -80,6 +80,17 @@ def _resolve_str_variable(variable, file_groups): for group in file_group.all_groups(): if any(ot == variable.points_to for ot in group.import_tokens): return group + + # Handle module paths like "infrastructure.http" matching file with import_tokens ["http"] + # This is needed for imports like "from infrastructure import http" + for file_group in file_groups: + for node in file_group.all_nodes(): + if any(variable.points_to.endswith('.' + ot) for ot in node.import_tokens): + return node + for group in file_group.all_groups(): + if any(variable.points_to.endswith('.' + ot) for ot in group.import_tokens): + return group + return OWNER_CONST.UNKNOWN_MODULE diff --git a/tests/test_code/py/module_prefix_imports/http.py b/tests/test_code/py/module_prefix_imports/http.py new file mode 100644 index 0000000..1d3c425 --- /dev/null +++ b/tests/test_code/py/module_prefix_imports/http.py @@ -0,0 +1,2 @@ +def request(method, url): + pass \ No newline at end of file diff --git a/tests/test_code/py/module_prefix_imports/main.py b/tests/test_code/py/module_prefix_imports/main.py new file mode 100644 index 0000000..5aa67e9 --- /dev/null +++ b/tests/test_code/py/module_prefix_imports/main.py @@ -0,0 +1,10 @@ +import repository.user_repository as user_repository +from repository.user_repository import get_users2 + + +def main(): + user_repository.get_users() + get_users2() + + +main() \ No newline at end of file diff --git a/tests/test_code/py/module_prefix_imports/user_repository.py b/tests/test_code/py/module_prefix_imports/user_repository.py new file mode 100644 index 0000000..3e1e981 --- /dev/null +++ b/tests/test_code/py/module_prefix_imports/user_repository.py @@ -0,0 +1,10 @@ +from infrastructure import http +from infrastructure.http import request + + +def get_users(): + http.request("GET", "https://api.example.com/users") + + +def get_users2(): + request("GET", "https://api.example.com/users") \ No newline at end of file diff --git a/tests/testdata.py b/tests/testdata.py index f3241a4..59a5355 100644 --- a/tests/testdata.py +++ b/tests/testdata.py @@ -554,6 +554,21 @@ "tzfile::(global)", "tzinfo::memorized_timedelta"] }, + { + "test_name": "module_prefix_imports", + "directory": "module_prefix_imports", + "comment": "Test that functions with module prefix are mapped correctly (issue #104)", + "expected_edges": [["main::(global)", "main::main"], + ["main::main", "user_repository::get_users"], + ["main::main", "user_repository::get_users2"], + ["user_repository::get_users", "http::request"], + ["user_repository::get_users2", "http::request"]], + "expected_nodes": ["http::request", + "main::(global)", + "main::main", + "user_repository::get_users", + "user_repository::get_users2"] + }, ], "js": [ {