Fix module prefix function mapping for Python imports

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
 
 

tests/test_code/py/module_prefix_imports/http.py

@@ -0,0 +1,2 @@
+def request(method, url):
+    pass

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()

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")

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": [
         {