Merge pull request #601 from Jaseci-Labs/real_python_raise

Make abs_path point to the real python file if possible
Jaseci-Labs · Sep 5, 2024 · 5882d8d · 5882d8d
2 parents 98017c4 + 7128bc2
commit 5882d8d
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 13 deletions.
diff --git a/jaclang/compiler/passes/main/py_collect_dep_pass.py b/jaclang/compiler/passes/main/py_collect_dep_pass.py
@@ -7,6 +7,8 @@
 from __future__ import annotations
 
 
+import os
+
 import jaclang.compiler.absyntree as ast
 from jaclang.compiler.passes import Pass
 from jaclang.settings import settings
@@ -28,18 +30,24 @@ def enter_node(self, node: ast.AstNode) -> None:
         if not isinstance(node, ast.AstSymbolNode):
             return
 
+        # Adding the path of the file related to the py import
         path: str = ""
         if isinstance(node, ast.ModulePath):
             if node.path:
                 path = ".".join([i.value for i in node.path])
             node.abs_path = self.ir.py_mod_dep_map.get(path)
+            if node.abs_path and os.path.isfile(node.abs_path.replace(".pyi", ".py")):
+                node.abs_path = node.abs_path.replace(".pyi", ".py")
+
         elif isinstance(node, ast.ModuleItem):
             imp = node.parent_of_type(ast.Import)
             mod_path_node = imp.get_all_sub_nodes(ast.ModulePath)[0]
             if mod_path_node.path:
                 path = ".".join([i.value for i in mod_path_node.path])
             path += f".{node.name.value}"
             node.abs_path = self.ir.py_mod_dep_map.get(path)
+            if node.abs_path and os.path.isfile(node.abs_path.replace(".pyi", ".py")):
+                node.abs_path = node.abs_path.replace(".pyi", ".py")
 
         if len(node.gen.mypy_ast) == 0:
             return

diff --git a/jaclang/compiler/passes/main/tests/fixtures/pygame_mock/__init__.pyi b/jaclang/compiler/passes/main/tests/fixtures/pygame_mock/__init__.pyi
@@ -0,0 +1,3 @@
+from .display import *
+from .color import *
+from .constants import *
diff --git a/jaclang/compiler/passes/main/tests/test_import_pass.py b/jaclang/compiler/passes/main/tests/test_import_pass.py
@@ -70,19 +70,19 @@ def test_py_raise_map(self) -> None:
         )
         assert isinstance(build.ir, ast.Module)
         p = {
-            "math": "jaclang/jaclang/vendor/mypy/typeshed/stdlib/math.pyi",
-            "pygame_mock": "pygame_mock/__init__.py",
-            "pygame_mock.color": "pygame_mock/color.py",
-            "pygame_mock.constants": "pygame_mock/constants.py",
-            "argparse": "jaclang/vendor/mypy/typeshed/stdlib/argparse.pyi",
-            "builtins": "jaclang/vendor/mypy/typeshed/stdlib/builtins.pyi",
-            "pygame_mock.display": "pygame_mock/display.py",
-            "os": "jaclang/vendor/mypy/typeshed/stdlib/os/__init__.pyi",
-            "genericpath": "jaclang/vendor/mypy/typeshed/stdlib/genericpath.pyi",
+            "math": r"jaclang/jaclang/vendor/mypy/typeshed/stdlib/math.pyi$",
+            "pygame_mock": r"pygame_mock/__init__.pyi$",
+            "pygame_mock.color": r"pygame_mock/color.py$",
+            "pygame_mock.constants": r"pygame_mock/constants.py$",
+            "argparse": r"jaclang/vendor/mypy/typeshed/stdlib/argparse.pyi$",
+            "builtins": r"jaclang/vendor/mypy/typeshed/stdlib/builtins.pyi$",
+            "pygame_mock.display": r"pygame_mock/display.py$",
+            "os": r"jaclang/vendor/mypy/typeshed/stdlib/os/__init__.pyi$",
+            "genericpath": r"jaclang/vendor/mypy/typeshed/stdlib/genericpath.pyi$",
         }
         for i in p:
             self.assertIn(i, build.ir.py_raise_map)
-            self.assertIn(p[i], re.sub(r".*fixtures/", "", build.ir.py_raise_map[i]))
+            self.assertRegex(re.sub(r".*fixtures/", "", build.ir.py_raise_map[i]), p[i])
 
     def test_py_raised_mods(self) -> None:
         """Basic test for pass."""

diff --git a/jaclang/langserve/engine.py b/jaclang/langserve/engine.py
@@ -349,7 +349,10 @@ def get_definition(
             elif node_selected.parent and isinstance(
                 node_selected.parent, ast.ModuleItem
             ):
-                path = node_selected.parent.from_mod_path.abs_path
+                path = (
+                    node_selected.parent.abs_path
+                    or node_selected.parent.from_mod_path.abs_path
+                )
                 try:  # TODO: Get rid of this when 'from' import is fixed
                     loc_range = tuple(
                         loc - 1 if loc > 0 else loc

diff --git a/jaclang/langserve/tests/test_server.py b/jaclang/langserve/tests/test_server.py
@@ -218,9 +218,9 @@ def test_go_to_definition_foolme(self) -> None:
         )
         lsp.deep_check(import_file)
         positions = [
-            (6, 39, "/pygame_mock/__init__.py:2:0-2:0"),
+            (6, 39, "/pygame_mock/__init__.pyi:2:0-2:0"),
             (6, 45, "/pygame_mock/constants.py:3:0-4:1"),
-            (7, 31, "/pygame_mock/__init__.py:2:0-2:0"),
+            (7, 31, "/pygame_mock/__init__.pyi:2:0-2:0"),
             (7, 35, "/pygame_mock/constants.py:3:0-4:1"),
             (20, 51, "/py_imp_test.jac:6:4-6:11"),
             (20, 64, "/pygame_mock/constants.py:4:3-4:15"),