Merge pull request #537 from Jaseci-Labs/mypy_symbol_resolve

Mypy symbol resolve

.pre-commit-config.yaml

@@ -16,7 +16,7 @@ repos:
     hooks:
       - id: flake8
         additional_dependencies: [pep8-naming, flake8_import_order, flake8_docstrings, flake8_comprehensions, flake8_bugbear, flake8_annotations, flake8_simplify]
-        exclude: "examples|vendor|langserve/tests"
+        exclude: "examples|vendor|langserve/tests|pygame_mock"
   - repo: https://github.com/pre-commit/mirrors-mypy
     rev: v1.10.1
     hooks:

jaclang/compiler/absyntree.py

@@ -630,8 +630,11 @@ def __init__(
         self.impl_mod: list[Module] = []
         self.test_mod: list[Module] = []
         self.mod_deps: dict[str, Module] = {}
+        self.py_mod_dep_map: dict[str, str] = {}
+        self.py_raise_map: dict[str, str] = {}
         self.registry = registry
         self.terminals: list[Token] = terminals
+        self.is_raised_from_py: bool = False
         AstNode.__init__(self, kid=kid)
         AstDocNode.__init__(self, doc=doc)
 

jaclang/compiler/passes/main/__init__.py

@@ -1,8 +1,8 @@
 """Collection of passes for Jac IR."""
 
 from .sub_node_tab_pass import SubNodeTabPass
-from .import_pass import JacImportPass, PyImportPass  # noqa: I100
 from .sym_tab_build_pass import SymTabBuildPass  # noqa: I100
+from .import_pass import JacImportPass, PyImportPass  # noqa: I100
 from .def_impl_match_pass import DeclImplMatchPass  # noqa: I100
 from .def_use_pass import DefUsePass  # noqa: I100
 from .pyout_pass import PyOutPass  # noqa: I100

jaclang/compiler/passes/main/access_modifier_pass.py

@@ -20,10 +20,20 @@ def after_pass(self) -> None:
         """After pass."""
         pass
 
-    def exit_node(self, node: ast.AstNode) -> None:
+    def exit_node(self, node: ast.AstNode) -> None:  # TODO: Move to debug pass
         """Exit node."""
         super().exit_node(node)
-        if settings.lsp_debug and isinstance(node, ast.NameAtom) and not node.sym:
+        if (
+            settings.lsp_debug
+            and isinstance(node, ast.NameAtom)
+            and not node.sym
+            and not node.parent_of_type(ast.Module).is_raised_from_py
+            and not (
+                node.sym_name == "py"
+                and node.parent
+                and isinstance(node.parent.parent, ast.Import)
+            )
+        ):
             self.warning(f"Name {node.sym_name} not present in symbol table")
 
     def access_check(self, node: ast.Name) -> None:

jaclang/compiler/passes/main/fuse_typeinfo_pass.py

@@ -10,7 +10,6 @@
 
 import jaclang.compiler.absyntree as ast
 from jaclang.compiler.passes import Pass
-from jaclang.compiler.symtable import SymbolTable
 from jaclang.settings import settings
 from jaclang.utils.helpers import pascal_to_snake
 from jaclang.vendor.mypy.nodes import Node as VNode  # bit of a hack
@@ -65,6 +64,35 @@ def __set_sym_table_link(self, node: ast.AstSymbolNode) -> None:
         if typ_sym_table != self.ir.sym_tab:
             node.name_spec.type_sym_tab = typ_sym_table
 
+    def __collect_python_dependencies(self, node: ast.AstNode) -> None:
+        assert isinstance(node, ast.AstSymbolNode)
+        assert isinstance(self.ir, ast.Module)
+
+        mypy_node = node.gen.mypy_ast[0]
+
+        if isinstance(mypy_node, MypyNodes.RefExpr):
+            node_full_name = mypy_node.node.fullname
+            if "." in node_full_name:
+                mod_name = node_full_name[: node_full_name.rindex(".")]
+            else:
+                mod_name = node_full_name
+
+            if mod_name not in self.ir.py_mod_dep_map:
+                self.__debug_print(
+                    f"Can't find a python file associated with {type(node)}::{node.loc}"
+                )
+                return
+
+            mode_path = self.ir.py_mod_dep_map[mod_name]
+            if mode_path.endswith(".jac"):
+                return
+
+            self.ir.py_raise_map[mod_name] = mode_path
+        else:
+            self.__debug_print(
+                f"Collect python dependencies is not supported in {type(node)}::{node.loc}"
+            )
+
     @staticmethod
     def __handle_node(
         func: Callable[[FuseTypeInfoPass, T], None]
@@ -86,6 +114,7 @@ def node_handler(self: FuseTypeInfoPass, node: T) -> None:
                 if len(node.gen.mypy_ast) == 1:
                     func(self, node)
                     self.__set_sym_table_link(node)
+                    self.__collect_python_dependencies(node)
 
                 # Jac node has multiple mypy nodes linked to it
                 elif len(node.gen.mypy_ast) > 1:
@@ -109,6 +138,7 @@ def node_handler(self: FuseTypeInfoPass, node: T) -> None:
                         )
                         func(self, node)
                         self.__set_sym_table_link(node)
+                        self.__collect_python_dependencies(node)
 
                 # Jac node doesn't have mypy nodes linked to it
                 else:
@@ -446,20 +476,16 @@ def exit_assignment(self, node: ast.Assignment) -> None:
                 if self_obj.type_sym_tab and isinstance(right_obj, ast.AstSymbolNode):
                     self_obj.type_sym_tab.def_insert(right_obj)
 
-    def exit_name(self, node: ast.Name) -> None:
-        """Add new symbols in the symbol table in case of atom trailer."""
-        if isinstance(node.parent, ast.AtomTrailer):
-            target_node = node.parent.target
-            if isinstance(target_node, ast.AstSymbolNode):
-                parent_symbol_table = target_node.type_sym_tab
-                if isinstance(parent_symbol_table, SymbolTable):
-                    node.sym = parent_symbol_table.lookup(node.sym_name)
-
-    # def exit_in_for_stmt(self, node: ast.InForStmt):
-    #     print(node.loc.mod_path, node.loc)
-    #     print(node.target, node.target.loc)
-    #     # node.sym_tab.def_insert()
-    #     # exit()
-
-    # def after_pass(self) -> None:
-    #     exit()
+    def exit_atom_trailer(self, node: ast.AtomTrailer) -> None:
+        """Adding symbol links to AtomTrailer right nodes."""
+        # This will fix adding the symbol links to nodes in atom trailer
+        # self.x.z = 5  # will add symbol links to both x and z
+        for i in range(1, len(node.as_attr_list)):
+            left = node.as_attr_list[i - 1]
+            right = node.as_attr_list[i]
+            # assert isinstance(left, ast.NameAtom)
+            # assert isinstance(right, ast.NameAtom)
+            if left.type_sym_tab and not isinstance(
+                right, ast.IndexSlice
+            ):  # TODO check why IndexSlice produce an issue
+                right.name_spec.sym = left.type_sym_tab.lookup(right.sym_name)

jaclang/compiler/passes/main/import_pass.py

@@ -6,19 +6,16 @@
 """
 
 import ast as py_ast
-import importlib.util
 import os
-import sys
 from typing import Optional
 
 
 import jaclang.compiler.absyntree as ast
 from jaclang.compiler.passes import Pass
-from jaclang.compiler.passes.main import SubNodeTabPass
-from jaclang.settings import settings
-from jaclang.utils.helpers import is_standard_lib_module
+from jaclang.compiler.passes.main import SubNodeTabPass, SymTabBuildPass
 from jaclang.utils.log import logging
 
+
 logger = logging.getLogger(__name__)
 
 
@@ -28,37 +25,29 @@ class JacImportPass(Pass):
     def before_pass(self) -> None:
         """Run once before pass."""
         self.import_table: dict[str, ast.Module] = {}
-        self.__py_imports: set[str] = set()
-        self.py_resolve_list: set[str] = set()
 
     def enter_module(self, node: ast.Module) -> None:
         """Run Importer."""
         self.cur_node = node
         self.import_table[node.loc.mod_path] = node
-        self.__py_imports.clear()
-        self.__annex_impl(node)
+        self.annex_impl(node)
         self.terminate()  # Turns off auto traversal for deliberate traversal
         self.run_again = True
         while self.run_again:
             self.run_again = False
             all_imports = self.get_all_sub_nodes(node, ast.ModulePath)
             for i in all_imports:
-                self.process_import(node, i)
+                self.process_import(i)
                 self.enter_module_path(i)
             SubNodeTabPass(prior=self, input_ir=node)
 
-        for i in self.get_all_sub_nodes(node, ast.AtomTrailer):
-            self.enter_atom_trailer(i)
-
-        node.mod_deps = self.import_table
+        node.mod_deps.update(self.import_table)
 
-    def process_import(self, node: ast.Module, i: ast.ModulePath) -> None:
+    def process_import(self, i: ast.ModulePath) -> None:
         """Process an import."""
         imp_node = i.parent_of_type(ast.Import)
         if imp_node.is_jac and not i.sub_module:
             self.import_jac_module(node=i)
-        elif imp_node.is_py:
-            self.__py_imports.add(i.path_str)
 
     def attach_mod_to_node(
         self, node: ast.ModulePath | ast.ModuleItem, mod: ast.Module | None
@@ -67,10 +56,11 @@ def attach_mod_to_node(
         if mod:
             self.run_again = True
             node.sub_module = mod
-            self.__annex_impl(mod)
+            self.annex_impl(mod)
             node.add_kids_right([mod], pos_update=False)
+            mod.parent = node
 
-    def __annex_impl(self, node: ast.Module) -> None:
+    def annex_impl(self, node: ast.Module) -> None:
         """Annex impl and test modules."""
         if node.stub_only:
             return
@@ -131,11 +121,6 @@ def enter_module_path(self, node: ast.ModulePath) -> None:
             node.sub_module.name = node.alias.value
         # Items matched during def/decl pass
 
-    def enter_atom_trailer(self, node: ast.AtomTrailer) -> None:
-        """Iterate on AtomTrailer nodes to get python paths to resolve."""
-        if node.as_attr_list[0].sym_name in self.__py_imports:
-            self.py_resolve_list.add(".".join([i.sym_name for i in node.as_attr_list]))
-
     def import_jac_module(self, node: ast.ModulePath) -> None:
         """Import a module."""
         self.cur_node = node  # impacts error reporting
@@ -214,58 +199,102 @@ class PyImportPass(JacImportPass):
     def before_pass(self) -> None:
         """Only run pass if settings are set to raise python."""
         super().before_pass()
-        if not settings.py_raise:
-            self.terminate()
 
-    def process_import(self, node: ast.Module, i: ast.ModulePath) -> None:
+    def __get_current_module(self, node: ast.AstNode) -> str:
+        parent = node.find_parent_of_type(ast.Module)
+        mod_list = []
+        while parent is not None:
+            mod_list.append(parent)
+            parent = parent.find_parent_of_type(ast.Module)
+        mod_list.reverse()
+        return ".".join(p.name for p in mod_list)
+
+    def process_import(self, i: ast.ModulePath) -> None:
         """Process an import."""
+        # Process import is orginally implemented to handle ModulePath in Jaclang
+        # This won't work with py imports as this will fail to import stuff in form of
+        #      from a import b
+        #      from a import (c, d, e)
+        # Solution to that is to get the import node and check the from loc then
+        # handle it based on if there a from loc or not
         imp_node = i.parent_of_type(ast.Import)
-        if (
-            imp_node.is_py
-            and not i.sub_module
-            and not is_standard_lib_module(i.path_str)
-        ):
-            mod = self.import_py_module(node=i, mod_path=node.loc.mod_path)
-            if mod:
-                i.sub_module = mod
-                i.add_kids_right([mod], pos_update=False)
-                if settings.py_raise_deep:
-                    self.run_again = True
+        if imp_node.is_py and not i.sub_module:
+            if imp_node.from_loc:
+                for j in imp_node.items.items:
+                    assert isinstance(j, ast.ModuleItem)
+                    mod_path = f"{imp_node.from_loc.path_str}.{j.name.sym_name}"
+                    self.import_py_module(
+                        parent_node=j,
+                        mod_path=mod_path,
+                        imported_mod_name=(
+                            j.name.sym_name if not j.alias else j.alias.sym_name
+                        ),
+                    )
+            else:
+                for j in imp_node.items.items:
+                    assert isinstance(j, ast.ModulePath)
+                    self.import_py_module(
+                        parent_node=j,
+                        mod_path=j.path_str,
+                        imported_mod_name=(
+                            j.path_str.replace(".", "")
+                            if not j.alias
+                            else j.alias.sym_name
+                        ),
+                    )
 
     def import_py_module(
-        self, node: ast.ModulePath, mod_path: str
+        self,
+        parent_node: ast.ModulePath | ast.ModuleItem,
+        imported_mod_name: str,
+        mod_path: str,
     ) -> Optional[ast.Module]:
         """Import a module."""
         from jaclang.compiler.passes.main import PyastBuildPass
 
-        base_dir = os.path.dirname(mod_path)
-        sys.path.append(base_dir)
+        assert isinstance(self.ir, ast.Module)
+
+        python_raise_map = self.ir.py_raise_map
+        file_to_raise = None
+
+        if mod_path in python_raise_map:
+            file_to_raise = python_raise_map[mod_path]
+        else:
+            resolved_mod_path = (
+                f"{self.__get_current_module(parent_node)}.{imported_mod_name}"
+            )
+            assert isinstance(self.ir, ast.Module)
+            resolved_mod_path = resolved_mod_path.replace(f"{self.ir.name}.", "")
+            file_to_raise = python_raise_map.get(resolved_mod_path)
+
+        if file_to_raise is None:
+            return None
+
         try:
-            # Dynamically import the module
-            spec = importlib.util.find_spec(node.path_str)
-            sys.path.remove(base_dir)
-            if spec and spec.origin and spec.origin not in {None, "built-in", "frozen"}:
-                if spec.origin in self.import_table:
-                    return self.import_table[spec.origin]
-                with open(spec.origin, "r", encoding="utf-8") as f:
+            if file_to_raise not in {None, "built-in", "frozen"}:
+                if file_to_raise in self.import_table:
+                    return self.import_table[file_to_raise]
+
+                with open(file_to_raise, "r", encoding="utf-8") as f:
                     mod = PyastBuildPass(
                         input_ir=ast.PythonModuleAst(
-                            py_ast.parse(f.read()), mod_path=spec.origin
+                            py_ast.parse(f.read()), mod_path=file_to_raise
                         ),
                     ).ir
+                    SubNodeTabPass(input_ir=mod, prior=self)
                 if mod:
-                    self.import_table[spec.origin] = mod
+                    mod.name = imported_mod_name
+                    self.import_table[file_to_raise] = mod
+                    self.attach_mod_to_node(parent_node, mod)
+                    SymTabBuildPass(input_ir=mod, prior=self)
                     return mod
                 else:
-                    raise self.ice(
-                        f"Failed to import python module {node.path_str}: {spec.origin}"
-                    )
+                    raise self.ice(f"Failed to import python module {mod_path}")
         except Exception as e:
-            if "Empty kid for Token ModulePath" in str(e) or "utf-8" in str(e):  # FIXME
-                return None
-            self.error(
-                f"Failed to import python module {node.path_str}: {e}",
-                node_override=node,
-            )
+            self.error(f"Failed to import python module {mod_path}")
             raise e
         return None
+
+    def annex_impl(self, node: ast.Module) -> None:
+        """Annex impl and test modules."""
+        return None

jaclang/compiler/passes/main/pyast_load_pass.py

@@ -128,6 +128,7 @@ class Module(mod):
             is_imported=False,
             kid=valid,
         )
+        ret.is_raised_from_py = True
         return self.nu(ret)
 
     def proc_function_def(