jaclang 0.7.2__py3-none-any.whl → 0.7.8__py3-none-any.whl

jaclang/langserve/tests/fixtures/circle_err.jac

@@ -59,15 +59,15 @@ print(f"Area of a {c.shape_type.value} with radius {RAD} using class: {c.area()}
 glob expected_area = 78.53981633974483;
 
 test calc_area   {
-    check.assertAlmostEqual(calculate_area(RAD), expected_area);
+    check assertAlmostEqual(calculate_area(RAD), expected_area);
 }
 
 test circle_area   {
     c = Circle(RAD);
-    check.assertAlmostEqual(c.area(), expected_area);
+    check assertAlmostEqual(c.area(), expected_area);
 }
 
 test circle_type   {
     c = Circle(RAD);
-    check.assertEqual(c.shape_type, ShapeType.CIRCLE);
+    check assertEqual(c.shape_type, ShapeType.CIRCLE);
 }

jaclang/langserve/tests/fixtures/circle_pure.test.jac

@@ -0,0 +1,15 @@
+glob expected_area = 78.53981633974483;
+
+test a1 {
+    check assertAlmostEqual(calculate_area(RAD), expected_area);
+}
+
+test a2 {
+    c = Circle(RAD);
+    check assertAlmostEqual(c.area(), expected_area);
+}
+
+test a3 {
+    c = Circle(RAD);
+    check assertEqual(c.shape_type, ShapeType.CIRCLE);
+}

jaclang/langserve/tests/fixtures/import_include_statements.jac

@@ -0,0 +1,6 @@
+import:py os;
+import:py from math, sqrt as square_root;
+import:py datetime as dt;
+import:jac from base_module_structure, add as adsd, subtract,x,Color as clr;
+import:jac base_module_structure as base_module_structure;
+import:py from py_import,add1 as ss, sub1 as subtract1,apple,Orange1;

jaclang/langserve/tests/fixtures/py_import.py

@@ -0,0 +1,26 @@
+"""This is a helper file."""
+
+
+def add1():
+    return 1 + 1
+
+
+def sub1():
+    return 1 - 1
+
+
+# classes
+class Orange1:
+    def __init__(self):
+        self.orange = 1
+
+    def get_orange(self):
+        return self.orange
+
+    def set_orange(self, orange):
+        self.orange = orange
+
+
+# variables
+orange2 = 1
+apple = 1

jaclang/langserve/tests/test_server.py

@@ -63,10 +63,10 @@ class TestJacLangServer(TestCase):
         self.assertEqual(len(lsp.modules), 1)
         self.assertEqual(lsp.modules[circle_file].diagnostics[0].range.start.line, 22)
         lsp.deep_check(circle_file)
-        self.assertEqual(len(lsp.modules), 1)
+        # self.assertEqual(len(lsp.modules), 1)
         self.assertEqual(lsp.modules[circle_file].diagnostics[0].range.start.line, 22)
         lsp.type_check(circle_file)
-        self.assertEqual(len(lsp.modules), 1)
+        # self.assertEqual(len(lsp.modules), 1)
         self.assertEqual(lsp.modules[circle_file].diagnostics[0].range.start.line, 22)
 
     def test_impl_stay_connected(self) -> None:
@@ -88,7 +88,7 @@ class TestJacLangServer(TestCase):
             "Circle class inherits from Shape.",
             lsp.get_hover_info(circle_file, pos).contents.value,
         )
-        lsp.type_check(circle_impl_file, force=True)
+        lsp.type_check(circle_impl_file)
         pos = lspt.Position(8, 11)
         self.assertIn(
             "ability) calculate_area: float",
@@ -105,9 +105,9 @@ class TestJacLangServer(TestCase):
         circle_impl_file = uris.from_fs_path(
             self.fixture_abs_path("circle_pure.impl.jac")
         )
-        lsp.quick_check(circle_impl_file, force=True)
-        lsp.deep_check(circle_impl_file, force=True)
-        lsp.type_check(circle_impl_file, force=True)
+        lsp.quick_check(circle_impl_file)
+        lsp.deep_check(circle_impl_file)
+        lsp.type_check(circle_impl_file)
         pos = lspt.Position(8, 11)
         self.assertIn(
             "ability) calculate_area: float",
@@ -143,16 +143,7 @@ class TestJacLangServer(TestCase):
         lsp.quick_check(circle_file)
         lsp.deep_check(circle_file)
         lsp.type_check(circle_file)
-        expected_string = (
-            "DocumentSymbol(name='calculate_area', kind=<SymbolKind.Function: 12>, range=9:0-9:43, "
-            "selection_range=9:0-9:43, detail=None, tags=None, deprecated=None, children=["
-            "DocumentSymbol(name='radius', kind=<SymbolKind.Variable: 13>, range=9:1-9:14, "
-            "selection_range=9:1-9:14, detail=None, tags=None, deprecated=None, children=[])])"
-        )
-        self.assertEqual(
-            expected_string, str((lsp.get_document_symbols(circle_file))[6])
-        )
-        self.assertEqual(10, len(lsp.get_document_symbols(circle_file)))
+        self.assertEqual(8, len(lsp.get_document_symbols(circle_file)))
 
     def test_go_to_definition(self) -> None:
         """Test that the go to definition is correct."""
@@ -169,7 +160,7 @@ class TestJacLangServer(TestCase):
             str(lsp.get_definition(circle_file, lspt.Position(9, 16))),
         )
         self.assertIn(
-            "fixtures/circle_pure.jac:12:0-17:1",
+            "fixtures/circle_pure.jac:13:11-13:16",
             str(lsp.get_definition(circle_file, lspt.Position(20, 17))),
         )
 
@@ -186,6 +177,90 @@ class TestJacLangServer(TestCase):
         lsp.deep_check(guess_game_file)
         lsp.type_check(guess_game_file)
         self.assertIn(
-            "guess_game4.jac:27:4-27:34",
+            "guess_game4.jac:27:8-27:21",
             str(lsp.get_definition(guess_game_file, lspt.Position(46, 45))),
         )
+
+    def test_test_annex(self) -> None:
+        """Test that the server doesn't run if there is a syntax error."""
+        lsp = JacLangServer()
+        # Set up the workspace path to "fixtures/"
+        workspace_path = self.fixture_abs_path("")
+        workspace = Workspace(workspace_path, lsp)
+        lsp.lsp._workspace = workspace
+        circle_file = uris.from_fs_path(self.fixture_abs_path("circle_pure.test.jac"))
+        lsp.quick_check(circle_file)
+        lsp.deep_check(circle_file)
+        lsp.type_check(circle_file)
+        pos = lspt.Position(13, 29)
+        self.assertIn(
+            "shape_type: circle_pure.ShapeType",
+            lsp.get_hover_info(circle_file, pos).contents.value,
+        )
+
+    def test_go_to_defintion_import(self) -> None:
+        """Test that the go to definition is correct."""
+        lsp = JacLangServer()
+        workspace_path = self.fixture_abs_path("")
+        workspace = Workspace(workspace_path, lsp)
+        lsp.lsp._workspace = workspace
+        import_file = uris.from_fs_path(
+            self.fixture_abs_path("import_include_statements.jac")
+        )
+        lsp.quick_check(import_file)
+        lsp.deep_check(import_file)
+        lsp.type_check(import_file)
+        positions = [
+            (2, 16, "datetime.py:0:0-0:0"),
+            (3, 17, "base_module_structure.jac:0:0-0:0"),
+            (3, 74, "base_module_structure.jac:23:0-23:5"),
+            (5, 65, "py_import.py:12:0-20:5"),
+            (5, 35, "py_import.py:3:0-4:5"),
+        ]
+
+        for line, char, expected in positions:
+            with self.subTest(line=line, char=char):
+                self.assertIn(
+                    expected,
+                    str(lsp.get_definition(import_file, lspt.Position(line, char))),
+                )
+
+    def test_sem_tokens(self) -> None:
+        """Test that the Semantic Tokens are generated correctly."""
+        lsp = JacLangServer()
+        workspace_path = self.fixture_abs_path("")
+        workspace = Workspace(workspace_path, lsp)
+        lsp.lsp._workspace = workspace
+        circle_file = uris.from_fs_path(self.fixture_abs_path("circle.jac"))
+        lsp.quick_check(circle_file)
+        lsp.deep_check(circle_file)
+        lsp.type_check(circle_file)
+        sem_list = lsp.get_semantic_tokens(circle_file).data
+        expected_counts = [
+            ("<JacSemTokenType.VARIABLE: 8>, <JacSemTokenModifier.READONLY: 4>", 206),
+            (
+                "<JacSemTokenType.PROPERTY: 9>, <JacSemTokenModifier.DEFINITION: 2>,",
+                112,
+            ),
+            (
+                "<JacSemTokenType.PARAMETER: 7>, <JacSemTokenModifier.DECLARATION: 1>,",
+                56,
+            ),
+            (
+                "<JacSemTokenType.FUNCTION: 12>, <JacSemTokenModifier.DECLARATION: 1>,",
+                25,
+            ),
+            ("<JacSemTokenType.METHOD: 13>, <JacSemTokenModifier.DECLARATION: 1>", 12),
+            ("<JacSemTokenType.ENUM: 3>, <JacSemTokenModifier.DECLARATION: 1>,", 37),
+            ("<JacSemTokenType.CLASS: 2>, <JacSemTokenModifier.DECLARATION: ", 162),
+            (
+                "<JacSemTokenType.NAMESPACE: 0>, <JacSemTokenModifier.DEFINITION: 2>,",
+                10,
+            ),
+            ("0, 0, 4,", 22),
+            ("0, 0, 3,", 192),
+            ("0, 0, 6, ", 65),
+            (" 0, 7, 3,", 3),
+        ]
+        for token_type, expected_count in expected_counts:
+            self.assertEqual(str(sem_list).count(token_type), expected_count)

jaclang/langserve/utils.py

@@ -2,12 +2,16 @@
 
 import asyncio
 import builtins
+import importlib.util
+import os
+import sys
 from functools import wraps
 from typing import Any, Awaitable, Callable, Coroutine, Optional, ParamSpec, TypeVar
 
 import jaclang.compiler.absyntree as ast
 from jaclang.compiler.codeloc import CodeLocInfo
 from jaclang.compiler.symtable import Symbol, SymbolTable
+from jaclang.utils.helpers import import_target_to_relative_path
 
 import lsprotocol.types as lspt
 
@@ -68,7 +72,7 @@ def find_deepest_symbol_node_at_pos(
         if isinstance(node, ast.AstSymbolNode):
             last_symbol_node = node
 
-        for child in node.kid:
+        for child in [i for i in node.kid if i.loc.mod_path == node.loc.mod_path]:
             if position_within_node(child, line, character):
                 deeper_node = find_deepest_symbol_node_at_pos(child, line, character)
                 if deeper_node is not None:
@@ -106,23 +110,27 @@ def collect_symbols(node: SymbolTable) -> list[lspt.DocumentSymbol]:
         return symbols
 
     for key, item in node.tab.items():
-        if key in dir(builtins):
-            continue
-        if item in [owner_sym(tab) for tab in node.kid]:
+        if (
+            key in dir(builtins)
+            or item in [owner_sym(tab) for tab in node.kid]
+            or item.decl.loc.mod_path != node.owner.loc.mod_path
+        ):
             continue
         else:
 
-            pos = create_range(item.defn[0].loc)
+            pos = create_range(item.decl.loc)
             symbol = lspt.DocumentSymbol(
                 name=key,
-                kind=kind_map(item.defn[0]),
+                kind=kind_map(item.decl),
                 range=pos,
                 selection_range=pos,
                 children=[],
             )
             symbols.append(symbol)
 
-    for sub_tab in node.kid:
+    for sub_tab in [
+        i for i in node.kid if i.owner.loc.mod_path == node.owner.loc.mod_path
+    ]:
         sub_symbols = collect_symbols(sub_tab)
 
         if isinstance(
@@ -146,7 +154,7 @@ def collect_symbols(node: SymbolTable) -> list[lspt.DocumentSymbol]:
 
 def owner_sym(table: SymbolTable) -> Optional[Symbol]:
     """Get owner sym."""
-    if table.has_parent() and isinstance(table.owner, ast.AstSymbolNode):
+    if table.parent and isinstance(table.owner, ast.AstSymbolNode):
         return table.parent.lookup(table.owner.sym_name)
     return None
 
@@ -154,8 +162,14 @@ def owner_sym(table: SymbolTable) -> Optional[Symbol]:
 def create_range(loc: CodeLocInfo) -> lspt.Range:
     """Create an lspt.Range from a location object."""
     return lspt.Range(
-        start=lspt.Position(line=loc.first_line - 1, character=loc.col_start - 1),
-        end=lspt.Position(line=loc.last_line - 1, character=loc.col_end - 1),
+        start=lspt.Position(
+            line=loc.first_line - 1 if loc.first_line > 0 else 0,
+            character=loc.col_start - 1 if loc.col_start > 0 else 0,
+        ),
+        end=lspt.Position(
+            line=loc.last_line - 1 if loc.last_line > 0 else 0,
+            character=loc.col_end - 1 if loc.col_end > 0 else 0,
+        ),
     )
 
 
@@ -178,3 +192,103 @@ def kind_map(sub_tab: ast.AstNode) -> lspt.SymbolKind:
             )
         )
     )
+
+
+def get_mod_path(mod_path: ast.ModulePath, name_node: ast.Name) -> str | None:
+    """Get path for a module import name."""
+    ret_target = None
+    module_location_path = mod_path.loc.mod_path
+    if mod_path.parent and (
+        (
+            isinstance(mod_path.parent.parent, ast.Import)
+            and mod_path.parent.parent.hint.tag.value == "py"
+        )
+        or (
+            isinstance(mod_path.parent, ast.Import)
+            and mod_path.parent.from_loc
+            and mod_path.parent.hint.tag.value == "py"
+        )
+    ):
+        if mod_path.path and name_node in mod_path.path:
+            temporary_path_str = ("." * mod_path.level) + ".".join(
+                [p.value for p in mod_path.path[: mod_path.path.index(name_node) + 1]]
+                if mod_path.path
+                else ""
+            )
+        else:
+            temporary_path_str = mod_path.path_str
+        sys.path.append(os.path.dirname(module_location_path))
+        spec = importlib.util.find_spec(temporary_path_str)
+        sys.path.remove(os.path.dirname(module_location_path))
+        if spec and spec.origin and spec.origin.endswith(".py"):
+            ret_target = spec.origin
+    elif mod_path.parent and (
+        (
+            isinstance(mod_path.parent.parent, ast.Import)
+            and mod_path.parent.parent.hint.tag.value == "jac"
+        )
+        or (
+            isinstance(mod_path.parent, ast.Import)
+            and mod_path.parent.from_loc
+            and mod_path.parent.hint.tag.value == "jac"
+        )
+    ):
+        ret_target = import_target_to_relative_path(
+            level=mod_path.level,
+            target=mod_path.path_str,
+            base_path=os.path.dirname(module_location_path),
+        )
+    return ret_target
+
+
+def get_item_path(mod_item: ast.ModuleItem) -> tuple[str, tuple[int, int]] | None:
+    """Get path."""
+    item_name = mod_item.name.value
+    if mod_item.from_parent.hint.tag.value == "py" and mod_item.from_parent.from_loc:
+        path = get_mod_path(mod_item.from_parent.from_loc, mod_item.name)
+        if path:
+            return get_definition_range(path, item_name)
+    elif mod_item.from_parent.hint.tag.value == "jac":
+        mod_node = mod_item.from_mod_path
+        if mod_node.sub_module and mod_node.sub_module._sym_tab:
+            for symbol_name, symbol in mod_node.sub_module._sym_tab.tab.items():
+                if symbol_name == item_name:
+                    return symbol.decl.loc.mod_path, (
+                        symbol.decl.loc.first_line - 1,
+                        symbol.decl.loc.last_line - 1,
+                    )
+    return None
+
+
+def get_definition_range(
+    filename: str, name: str
+) -> tuple[str, tuple[int, int]] | None:
+    """Get the start and end line of a function or class definition in a file."""
+    import ast
+
+    with open(filename, "r") as file:
+        source = file.read()
+
+    tree = ast.parse(source)
+
+    for node in ast.walk(tree):
+        if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name == name:
+            start_line = node.lineno
+            end_line = (
+                node.body[-1].end_lineno
+                if hasattr(node.body[-1], "end_lineno")
+                else node.body[-1].lineno
+            )
+            if start_line and end_line:
+                return filename, (start_line - 1, end_line - 1)
+        elif isinstance(node, ast.Assign):
+            for target in node.targets:
+                if isinstance(target, ast.Name) and target.id == name:
+                    start_line = node.lineno
+                    end_line = (
+                        node.end_lineno if hasattr(node, "end_lineno") else node.lineno
+                    )
+                    if start_line and end_line:
+                        return filename, (start_line - 1, end_line - 1)
+
+    return None

jaclang/plugin/builtin.py

@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from typing import Optional
-    from jaclang.core.construct import NodeArchitype
+    from jaclang.core.constructs import NodeArchitype
 
 
 def dotgen(

jaclang/plugin/default.py

@@ -14,14 +14,8 @@ from typing import Any, Callable, Optional, Type, Union
 
 from jaclang.compiler.absyntree import Module
 from jaclang.compiler.constant import EdgeDir, colors
-from jaclang.core.aott import (
-    aott_raise,
-    extract_non_primary_type,
-    get_all_type_explanations,
-    get_info_types,
-    get_input_information,
-)
-from jaclang.core.construct import (
+from jaclang.compiler.semtable import SemInfo, SemRegistry, SemScope
+from jaclang.core.constructs import (
     Architype,
     DSFunc,
     EdgeAnchor,
@@ -39,7 +33,6 @@ from jaclang.core.construct import (
     exec_context,
 )
 from jaclang.core.importer import jac_importer
-from jaclang.core.registry import SemInfo, SemRegistry, SemScope
 from jaclang.core.utils import traverse_graph
 from jaclang.plugin.feature import JacFeature as Jac
 from jaclang.plugin.spec import T
@@ -110,7 +103,9 @@ class JacFeatureDefaults:
         """Create a new architype."""
         for i in on_entry + on_exit:
             i.resolve(cls)
-        if not issubclass(cls, arch_base):
+        if not hasattr(cls, "_jac_entry_funcs_") or not hasattr(
+            cls, "_jac_exit_funcs_"
+        ):
             # Saving the module path and reassign it after creating cls
             # So the jac modules are part of the correct module
             cur_module = cls.__module__
@@ -216,7 +211,7 @@ class JacFeatureDefaults:
         cachable: bool,
         mdl_alias: Optional[str],
         override_name: Optional[str],
-        mod_bundle: Optional[Module],
+        mod_bundle: Optional[Module | str],
         lng: Optional[str],
         items: Optional[dict[str, Union[str, bool]]],
     ) -> Optional[types.ModuleType]:
@@ -335,7 +330,13 @@ class JacFeatureDefaults:
     @hookimpl
     def ignore(
         walker: WalkerArchitype,
-        expr: list[NodeArchitype | EdgeArchitype] | NodeArchitype | EdgeArchitype,
+        expr: (
+            list[NodeArchitype | EdgeArchitype]
+            | list[NodeArchitype]
+            | list[EdgeArchitype]
+            | NodeArchitype
+            | EdgeArchitype
+        ),
     ) -> bool:
         """Jac's ignore stmt feature."""
         return walker._jac_.ignore_node(expr)
@@ -344,7 +345,13 @@ class JacFeatureDefaults:
     @hookimpl
     def visit_node(
         walker: WalkerArchitype,
-        expr: list[NodeArchitype | EdgeArchitype] | NodeArchitype | EdgeArchitype,
+        expr: (
+            list[NodeArchitype | EdgeArchitype]
+            | list[NodeArchitype]
+            | list[EdgeArchitype]
+            | NodeArchitype
+            | EdgeArchitype
+        ),
     ) -> bool:
         """Jac's visit stmt feature."""
         if isinstance(walker, WalkerArchitype):
@@ -481,14 +488,14 @@ class JacFeatureDefaults:
     @hookimpl
     def build_edge(
         is_undirected: bool,
-        conn_type: Optional[Type[EdgeArchitype]],
+        conn_type: Optional[Type[EdgeArchitype] | EdgeArchitype],
         conn_assign: Optional[tuple[tuple, tuple]],
     ) -> Callable[[], EdgeArchitype]:
         """Jac's root getter."""
         conn_type = conn_type if conn_type else GenericEdge
 
         def builder() -> EdgeArchitype:
-            edge = conn_type()
+            edge = conn_type() if isinstance(conn_type, type) else conn_type
             edge._jac_.is_undirected = is_undirected
             if conn_assign:
                 for fld, val in zip(conn_assign[0], conn_assign[1]):
@@ -602,7 +609,7 @@ class JacFeatureDefaults:
                     )
         if isinstance(attr_sem_info, SemInfo) and isinstance(attr_scope, SemScope):
             return attr_sem_info.semstr, attr_scope.as_type_str
-        return "", ""  #
+        return "", ""
 
     @staticmethod
     @hookimpl
@@ -618,68 +625,9 @@ class JacFeatureDefaults:
         action: str,
     ) -> Any:  # noqa: ANN401
         """Jac's with_llm feature."""
-        with open(
-            os.path.join(
-                os.path.dirname(file_loc),
-                "__jac_gen__",
-                os.path.basename(file_loc).replace(".jac", ".registry.pkl"),
-            ),
-            "rb",
-        ) as f:
-            mod_registry = pickle.load(f)
-
-        outputs = outputs[0] if isinstance(outputs, list) else outputs
-        _scope = SemScope.get_scope_from_str(scope)
-        assert _scope is not None
-
-        method = model_params.pop("method") if "method" in model_params else "Normal"
-        available_methods = model.MTLLM_METHOD_PROMPTS.keys()
-        assert (
-            method in available_methods
-        ), f"Invalid method: {method}. Select from {available_methods}"
-
-        context = (
-            "\n".join(model_params.pop("context")) if "context" in model_params else ""
-        )
-
-        type_collector: list = []
-        incl_info = [x for x in incl_info if not isinstance(x[1], type)]
-        information, collected_types = get_info_types(_scope, mod_registry, incl_info)
-        type_collector.extend(collected_types)
-
-        inputs_information = get_input_information(inputs, type_collector)
-
-        output_information = f"{outputs[0]} ({outputs[1]})"
-        type_collector.extend(extract_non_primary_type(outputs[1]))
-        output_type_explanations = "\n".join(
-            list(
-                get_all_type_explanations(
-                    extract_non_primary_type(outputs[1]), mod_registry
-                ).values()
-            )
-        )
-
-        type_explanations_list = list(
-            get_all_type_explanations(type_collector, mod_registry).values()
-        )
-        type_explanations = "\n".join(type_explanations_list)
-
-        meaning_out = aott_raise(
-            model=model,
-            information=information,
-            inputs_information=inputs_information,
-            output_information=output_information,
-            type_explanations=type_explanations,
-            action=action,
-            context=context,
-            method=method,
-            tools=[],
-            model_params=model_params,
-        )
-        output = model.resolve_output(
-            meaning_out, outputs[0], outputs[1], output_type_explanations
+        raise ImportError(
+            "mtllm is not installed. Please install it with `pip install mtllm`."
         )
-        return output
 
 
 class JacBuiltin:

jaclang/plugin/feature.py

@@ -3,10 +3,10 @@
 from __future__ import annotations
 
 import types
-from typing import Any, Callable, Optional, Type, Union
+from typing import Any, Callable, Optional, Type, TypeAlias, Union
 
 from jaclang.compiler.absyntree import Module
-from jaclang.core.construct import (
+from jaclang.core.constructs import (
     Architype,
     EdgeArchitype,
     Memory,
@@ -31,7 +31,13 @@ class JacFeature:
 
     import abc
     from jaclang.compiler.constant import EdgeDir
-    from jaclang.plugin.spec import DSFunc
+    from jaclang.core.constructs import DSFunc
+
+    RootType: TypeAlias = Root
+    Obj: TypeAlias = Architype
+    Node: TypeAlias = NodeArchitype
+    Edge: TypeAlias = EdgeArchitype
+    Walker: TypeAlias = WalkerArchitype
 
     @staticmethod
     def context(session: str = "") -> ExecutionContext:
@@ -96,7 +102,7 @@ class JacFeature:
         cachable: bool = True,
         mdl_alias: Optional[str] = None,
         override_name: Optional[str] = None,
-        mod_bundle: Optional[Module] = None,
+        mod_bundle: Optional[Module | str] = None,
         lng: Optional[str] = "jac",
         items: Optional[dict[str, Union[str, bool]]] = None,
     ) -> Optional[types.ModuleType]:
@@ -160,7 +166,13 @@ class JacFeature:
     @staticmethod
     def ignore(
         walker: WalkerArchitype,
-        expr: list[NodeArchitype | EdgeArchitype] | NodeArchitype | EdgeArchitype,
+        expr: (
+            list[NodeArchitype | EdgeArchitype]
+            | list[NodeArchitype]
+            | list[EdgeArchitype]
+            | NodeArchitype
+            | EdgeArchitype
+        ),
     ) -> bool:  # noqa: ANN401
         """Jac's ignore stmt feature."""
         return pm.hook.ignore(walker=walker, expr=expr)
@@ -168,7 +180,13 @@ class JacFeature:
     @staticmethod
     def visit_node(
         walker: WalkerArchitype,
-        expr: list[NodeArchitype | EdgeArchitype] | NodeArchitype | EdgeArchitype,
+        expr: (
+            list[NodeArchitype | EdgeArchitype]
+            | list[NodeArchitype]
+            | list[EdgeArchitype]
+            | NodeArchitype
+            | EdgeArchitype
+        ),
     ) -> bool:  # noqa: ANN401
         """Jac's visit stmt feature."""
         return pm.hook.visit_node(walker=walker, expr=expr)
@@ -245,7 +263,7 @@ class JacFeature:
     @staticmethod
     def build_edge(
         is_undirected: bool,
-        conn_type: Optional[Type[EdgeArchitype]],
+        conn_type: Optional[Type[EdgeArchitype] | EdgeArchitype],
         conn_assign: Optional[tuple[tuple, tuple]],
     ) -> Callable[[], EdgeArchitype]:
         """Jac's root getter."""