jaclang 0.8.1__py3-none-any.whl → 0.8.2__py3-none-any.whl

jaclang/compiler/passes/main/pyast_gen_pass.py

@@ -31,6 +31,63 @@ from jaclang.settings import settings
 
 T = TypeVar("T", bound=ast3.AST)
 
+# Mapping of Jac tokens to corresponding Python AST operator classes. This
+# helps keep the implementation of ``exit_token`` concise and easier to
+# maintain.
+TOKEN_AST_MAP: dict[Tok, type[ast3.AST]] = {
+    Tok.KW_AND: ast3.And,
+    Tok.KW_OR: ast3.Or,
+    Tok.PLUS: ast3.Add,
+    Tok.ADD_EQ: ast3.Add,
+    Tok.BW_AND: ast3.BitAnd,
+    Tok.BW_AND_EQ: ast3.BitAnd,
+    Tok.BW_OR: ast3.BitOr,
+    Tok.BW_OR_EQ: ast3.BitOr,
+    Tok.BW_XOR: ast3.BitXor,
+    Tok.BW_XOR_EQ: ast3.BitXor,
+    Tok.DIV: ast3.Div,
+    Tok.DIV_EQ: ast3.Div,
+    Tok.FLOOR_DIV: ast3.FloorDiv,
+    Tok.FLOOR_DIV_EQ: ast3.FloorDiv,
+    Tok.LSHIFT: ast3.LShift,
+    Tok.LSHIFT_EQ: ast3.LShift,
+    Tok.MOD: ast3.Mod,
+    Tok.MOD_EQ: ast3.Mod,
+    Tok.STAR_MUL: ast3.Mult,
+    Tok.MUL_EQ: ast3.Mult,
+    Tok.DECOR_OP: ast3.MatMult,
+    Tok.MATMUL_EQ: ast3.MatMult,
+    Tok.STAR_POW: ast3.Pow,
+    Tok.STAR_POW_EQ: ast3.Pow,
+    Tok.RSHIFT: ast3.RShift,
+    Tok.RSHIFT_EQ: ast3.RShift,
+    Tok.MINUS: ast3.Sub,
+    Tok.SUB_EQ: ast3.Sub,
+    Tok.BW_NOT: ast3.Invert,
+    Tok.BW_NOT_EQ: ast3.Invert,
+    Tok.NOT: ast3.Not,
+    Tok.EQ: ast3.NotEq,
+    Tok.EE: ast3.Eq,
+    Tok.GT: ast3.Gt,
+    Tok.GTE: ast3.GtE,
+    Tok.KW_IN: ast3.In,
+    Tok.KW_IS: ast3.Is,
+    Tok.KW_ISN: ast3.IsNot,
+    Tok.LT: ast3.Lt,
+    Tok.LTE: ast3.LtE,
+    Tok.NE: ast3.NotEq,
+    Tok.KW_NIN: ast3.NotIn,
+}
+
+# Mapping of unary operator tokens to their Python AST counterparts used in
+# ``exit_unary_expr``.
+UNARY_OP_MAP: dict[Tok, type[ast3.unaryop]] = {
+    Tok.NOT: ast3.Not,
+    Tok.BW_NOT: ast3.Invert,
+    Tok.PLUS: ast3.UAdd,
+    Tok.MINUS: ast3.USub,
+}
+
 
 class PyastGenPass(UniPass):
     """Jac blue transpilation to python pass."""
@@ -114,69 +171,55 @@ class PyastGenPass(UniPass):
             )
         )
 
-    def needs_typing(self) -> None:
-        """Check if enum is needed."""
-        if self.needs_typing.__name__ in self.already_added:
+    def _add_preamble_once(self, key: str, node: ast3.AST) -> None:
+        """Append an import statement to the preamble once."""
+        if key in self.already_added:
             return
-        self.preamble.append(
-            self.sync(
-                ast3.Import(
-                    names=[
-                        self.sync(
-                            ast3.alias(name="typing"),
-                            jac_node=self.ir_out,
-                        ),
-                    ]
-                ),
-                jac_node=self.ir_out,
-            )
+        self.preamble.append(self.sync(node, jac_node=self.ir_out))
+        self.already_added.append(key)
+
+    def needs_typing(self) -> None:
+        """Ensure typing is imported only once."""
+        self._add_preamble_once(
+            self.needs_typing.__name__,
+            ast3.Import(
+                names=[self.sync(ast3.alias(name="typing"), jac_node=self.ir_out)]
+            ),
         )
-        self.already_added.append(self.needs_typing.__name__)
 
     def needs_enum(self) -> None:
-        """Check if enum is needed."""
-        if self.needs_enum.__name__ in self.already_added:
-            return
-        self.preamble.append(
-            self.sync(
-                ast3.ImportFrom(
-                    module="enum",
-                    names=[
-                        self.sync(ast3.alias(name="Enum", asname=None)),
-                        self.sync(ast3.alias(name="auto", asname=None)),
-                    ],
-                    level=0,
-                ),
-                jac_node=self.ir_out,
-            )
+        """Ensure Enum utilities are imported only once."""
+        self._add_preamble_once(
+            self.needs_enum.__name__,
+            ast3.ImportFrom(
+                module="enum",
+                names=[
+                    self.sync(ast3.alias(name="Enum", asname=None)),
+                    self.sync(ast3.alias(name="auto", asname=None)),
+                ],
+                level=0,
+            ),
         )
-        self.already_added.append(self.needs_enum.__name__)
 
     def needs_future(self) -> None:
-        """Check if enum is needed."""
-        if self.needs_future.__name__ in self.already_added:
-            return
-        self.preamble.append(
-            self.sync(
-                ast3.ImportFrom(
-                    module="concurrent.futures",
-                    names=[
-                        self.sync(ast3.alias(name="Future", asname=None)),
-                    ],
-                    level=0,
-                ),
-                jac_node=self.ir_out,
-            )
+        """Ensure concurrent Future is imported only once."""
+        self._add_preamble_once(
+            self.needs_future.__name__,
+            ast3.ImportFrom(
+                module="concurrent.futures",
+                names=[self.sync(ast3.alias(name="Future", asname=None))],
+                level=0,
+            ),
         )
-        self.already_added.append(self.needs_future.__name__)
 
     def flatten(self, body: list[T | list[T] | None]) -> list[T]:
-        new_body = []
-        for i in body:
-            if isinstance(i, list):
-                new_body += i
-            elif i is not None:
-                new_body.append(i) if i else None
+        """Flatten a list of items or lists into a single list."""
+        new_body: list[T] = []
+        for item in body:
+            if isinstance(item, list):
+                new_body.extend(item)
+            elif item is not None:
+                new_body.append(item)
         return new_body
 
     def sync(
@@ -267,14 +310,10 @@ class PyastGenPass(UniPass):
         attr_node: ast3.Name | ast3.Attribute = self.sync(
             ast3.Name(id=attribute_list[0], ctx=ast3.Load()), sync_node_list[0]
         )
-        for i in range(len(attribute_list)):
-            if i == 0:
-                continue
+        for attr, sync_node in zip(attribute_list[1:], sync_node_list[1:]):
             attr_node = self.sync(
-                ast3.Attribute(
-                    value=attr_node, attr=attribute_list[i], ctx=ast3.Load()
-                ),
-                sync_node_list[i],
+                ast3.Attribute(value=attr_node, attr=attr, ctx=ast3.Load()),
+                sync_node,
             )
         return attr_node
 
@@ -426,40 +465,8 @@ class PyastGenPass(UniPass):
             )
 
     def exit_import(self, node: uni.Import) -> None:
-        path_alias: dict[str, Optional[str]] = (
-            {node.from_loc.dot_path_str: None} if node.from_loc else {}
-        )
-        imp_from = {}
-        if node.items:
-            for item in node.items:
-                if isinstance(item, uni.ModuleItem):
-                    imp_from[item.name.sym_name] = (
-                        item.alias.sym_name if item.alias else None
-                    )
-                elif isinstance(item, uni.ModulePath):
-                    path_alias[item.dot_path_str] = (
-                        item.alias.sym_name if item.alias else None
-                    )
-
-        item_names: list[ast3.expr] = []
-        item_keys: list[ast3.Constant] = []
-        item_values: list[ast3.Constant] = []
-        for k, v in imp_from.items():
-            item_keys.append(self.sync(ast3.Constant(value=k)))
-            item_values.append(self.sync(ast3.Constant(value=v)))
-            item_names.append(
-                self.sync(
-                    ast3.Name(
-                        id=v or k,
-                        ctx=ast3.Store(),
-                    )
-                )
-            )
-        path_named_value: str
+        """Exit import node."""
         py_nodes: list[ast3.AST] = []
-        typecheck_nodes: list[ast3.AST] = []
-        runtime_nodes: list[ast3.AST] = []
-
         if node.doc:
             py_nodes.append(
                 self.sync(
@@ -468,255 +475,28 @@ class PyastGenPass(UniPass):
                 )
             )
 
-        for path, alias in path_alias.items():
-            path_named_value = ("_jac_inc_" if node.is_absorb else "") + (
-                alias if alias else path
-            ).lstrip(".").split(".")[0]
-            # target_named_value = ""
-            # for i in path.split("."):
-            #     target_named_value += i if i else "."
-            #     if i:
-            #         break
-
-            args = [
-                self.sync(
-                    ast3.Constant(value=path),
-                ),
-                self.sync(
-                    ast3.Name(
-                        id="__file__",
-                        ctx=ast3.Load(),
-                    )
-                ),
-            ]
-            keywords = []
-
-            if node.is_absorb:
-                args.append(self.sync(ast3.Constant(value=node.is_absorb)))
-
-            if alias is not None:
-                keywords.append(
-                    self.sync(
-                        ast3.keyword(
-                            arg="mdl_alias",
-                            value=self.sync(
-                                ast3.Constant(value=alias),
-                            ),
-                        )
-                    )
-                )
-
-            if item_keys and item_values:
-                keywords.append(
-                    self.sync(
-                        ast3.keyword(
-                            arg="items",
-                            value=self.sync(
-                                ast3.Dict(
-                                    keys=cast(list[ast3.expr | None], item_keys),
-                                    values=cast(list[ast3.expr], item_values),
-                                ),
-                            ),
-                        )
-                    )
-                )
-
-            runtime_nodes.append(
-                self.sync(
-                    ast3.Assign(
-                        targets=(
-                            [
-                                self.sync(
-                                    ast3.Tuple(
-                                        elts=(
-                                            item_names
-                                            or [
-                                                self.sync(
-                                                    ast3.Name(
-                                                        id=path_named_value,
-                                                        ctx=ast3.Store(),
-                                                    )
-                                                )
-                                            ]
-                                        ),
-                                        ctx=ast3.Store(),
-                                    )
-                                )
-                            ]
-                        ),
-                        value=self.sync(
-                            ast3.Call(
-                                func=self.jaclib_obj("py_jac_import"),
-                                args=args,
-                                keywords=keywords,
-                            )
-                        ),
-                    ),
-                ),
-            )
-        if node.is_absorb:
-            absorb_exec = f"={path_named_value}.__dict__['"
-            runtime_nodes.append(
-                self.sync(
-                    ast3.For(
-                        target=self.sync(ast3.Name(id="i", ctx=ast3.Store())),
-                        iter=self.sync(
-                            ast3.IfExp(
-                                test=self.sync(
-                                    ast3.Compare(
-                                        left=self.sync(ast3.Constant(value="__all__")),
-                                        ops=[self.sync(ast3.In())],
-                                        comparators=[
-                                            self.sync(
-                                                ast3.Attribute(
-                                                    value=self.sync(
-                                                        ast3.Name(
-                                                            id=path_named_value,
-                                                            ctx=ast3.Load(),
-                                                        )
-                                                    ),
-                                                    attr="__dict__",
-                                                    ctx=ast3.Load(),
-                                                )
-                                            )
-                                        ],
-                                    )
-                                ),
-                                body=self.sync(
-                                    ast3.Attribute(
-                                        value=self.sync(
-                                            ast3.Name(
-                                                id=path_named_value, ctx=ast3.Load()
-                                            )
-                                        ),
-                                        attr="__all__",
-                                        ctx=ast3.Load(),
-                                    )
-                                ),
-                                orelse=self.sync(
-                                    ast3.Attribute(
-                                        value=self.sync(
-                                            ast3.Name(
-                                                id=path_named_value, ctx=ast3.Load()
-                                            )
-                                        ),
-                                        attr="__dict__",
-                                        ctx=ast3.Load(),
-                                    )
-                                ),
-                            )
-                        ),
-                        body=[
-                            self.sync(
-                                ast3.If(
-                                    test=self.sync(
-                                        ast3.UnaryOp(
-                                            op=self.sync(ast3.Not()),
-                                            operand=self.sync(
-                                                ast3.Call(
-                                                    func=self.sync(
-                                                        ast3.Attribute(
-                                                            value=self.sync(
-                                                                ast3.Name(
-                                                                    id="i",
-                                                                    ctx=ast3.Load(),
-                                                                )
-                                                            ),
-                                                            attr="startswith",
-                                                            ctx=ast3.Load(),
-                                                        )
-                                                    ),
-                                                    args=[
-                                                        self.sync(
-                                                            ast3.Constant(value="_")
-                                                        )
-                                                    ],
-                                                    keywords=[],
-                                                )
-                                            ),
-                                        )
-                                    ),
-                                    body=[
-                                        self.sync(
-                                            ast3.Expr(
-                                                value=self.sync(
-                                                    ast3.Call(
-                                                        func=self.sync(
-                                                            ast3.Name(
-                                                                id="exec",
-                                                                ctx=ast3.Load(),
-                                                            )
-                                                        ),
-                                                        args=[
-                                                            self.sync(
-                                                                ast3.JoinedStr(
-                                                                    values=[
-                                                                        self.sync(
-                                                                            ast3.FormattedValue(
-                                                                                value=self.sync(
-                                                                                    ast3.Name(
-                                                                                        id="i",
-                                                                                        ctx=ast3.Load(),
-                                                                                    )
-                                                                                ),
-                                                                                conversion=-1,
-                                                                            )
-                                                                        ),
-                                                                        self.sync(
-                                                                            ast3.Constant(
-                                                                                value=absorb_exec
-                                                                            )
-                                                                        ),
-                                                                        self.sync(
-                                                                            ast3.FormattedValue(
-                                                                                value=self.sync(
-                                                                                    ast3.Name(
-                                                                                        id="i",
-                                                                                        ctx=ast3.Load(),
-                                                                                    )
-                                                                                ),
-                                                                                conversion=-1,
-                                                                            )
-                                                                        ),
-                                                                        self.sync(
-                                                                            ast3.Constant(
-                                                                                value="']"
-                                                                            )
-                                                                        ),
-                                                                    ]
-                                                                )
-                                                            )
-                                                        ],
-                                                        keywords=[],
-                                                    )
-                                                )
-                                            )
-                                        )
-                                    ],
-                                    orelse=[],
-                                )
-                            )
-                        ],
-                        orelse=[],
-                    )
-                )
-            )
         if node.is_absorb:
+            # This is `include "module_name";` which becomes `from module_name import *`
             source = node.items[0]
             if not isinstance(source, uni.ModulePath):
                 raise self.ice()
-            typecheck_nodes.append(
+
+            module_name_parts = [p.value for p in source.path] if source.path else []
+            module_name = ".".join(module_name_parts) if module_name_parts else None
+
+            py_nodes.append(
                 self.sync(
                     py_node=ast3.ImportFrom(
-                        module=(source.dot_path_str.lstrip(".") if source else None),
+                        module=module_name,
                         names=[self.sync(ast3.alias(name="*"), node)],
-                        level=0,
+                        level=source.level,
                     ),
                     jac_node=node,
                 )
             )
         elif not node.from_loc:
-            typecheck_nodes.append(
+            # This is `import module1, module2 as alias;`
+            py_nodes.append(
                 self.sync(
                     ast3.Import(
                         names=[
@@ -728,32 +508,25 @@ class PyastGenPass(UniPass):
                 )
             )
         else:
-            typecheck_nodes.append(
+            # This is `from module import item1, item2 as alias;`
+            module_name_parts = (
+                [p.value for p in node.from_loc.path] if node.from_loc.path else []
+            )
+            module_name = ".".join(module_name_parts) if module_name_parts else None
+
+            py_nodes.append(
                 self.sync(
                     ast3.ImportFrom(
-                        module=(
-                            node.from_loc.dot_path_str.lstrip(".")
-                            if node.from_loc
-                            else None
-                        ),
+                        module=module_name,
                         names=[
                             cast(ast3.alias, i)
                             for item in node.items
                             for i in item.gen.py_ast
                         ],
-                        level=0,
+                        level=node.from_loc.level,
                     )
                 )
             )
-        py_nodes.append(
-            self.sync(
-                ast3.If(
-                    test=self.jaclib_obj("TYPE_CHECKING"),
-                    body=[cast(ast3.stmt, node) for node in typecheck_nodes],
-                    orelse=[cast(ast3.stmt, node) for node in runtime_nodes],
-                )
-            )
-        )
         node.gen.py_ast = py_nodes
 
     def exit_module_path(self, node: uni.ModulePath) -> None:
@@ -867,6 +640,17 @@ class PyastGenPass(UniPass):
         if isinstance(node.body, uni.ImplDef):
             self.traverse(node.body)
 
+    def gen_llm_call_override(self, node: uni.FuncCall) -> list[ast3.AST]:
+        """Generate python ast nodes for llm function body override syntax.
+
+        example:
+            foo() by llm();
+        """
+        # to Avoid circular import
+        from jaclang.runtimelib.machine import JacMachineInterface
+
+        return JacMachineInterface.gen_llm_call_override(self, node)
+
     def gen_llm_body(self, node: uni.Ability) -> list[ast3.AST]:
         """Generate the by LLM body."""
         # to Avoid circular import
@@ -988,11 +772,7 @@ class PyastGenPass(UniPass):
                     ),
                     body=[cast(ast3.stmt, i) for i in body],
                     decorator_list=[cast(ast3.expr, i) for i in decorator_list],
-                    returns=(
-                        cast(ast3.expr, node.signature.return_type.gen.py_ast[0])
-                        if node.signature and node.signature.return_type
-                        else self.sync(ast3.Constant(value=None))
-                    ),
+                    returns=self.sync(ast3.Constant(value=None)),
                     type_params=[],
                 )
             )
@@ -1829,25 +1609,9 @@ class PyastGenPass(UniPass):
         ]
 
     def exit_await_expr(self, node: uni.AwaitExpr) -> None:
-        parent_node = node.parent
-        while parent_node and (parent_node := parent_node.parent):
-            if hasattr(parent_node, "is_async") and parent_node.is_async:
-                node.gen.py_ast = [
-                    self.sync(
-                        ast3.Await(value=cast(ast3.expr, node.target.gen.py_ast[0]))
-                    )
-                ]
-                break
-        else:
-            node.gen.py_ast = [
-                self.sync(
-                    ast3.Call(
-                        func=self.jaclib_obj("await_obj"),
-                        args=[cast(ast3.expr, node.target.gen.py_ast[0])],
-                        keywords=[],
-                    )
-                )
-            ]
+        node.gen.py_ast = [
+            self.sync(ast3.Await(value=cast(ast3.expr, node.target.gen.py_ast[0])))
+        ]
 
     def exit_global_stmt(self, node: uni.GlobalStmt) -> None:
         py_nodes = []
@@ -2180,42 +1944,17 @@ class PyastGenPass(UniPass):
         ]
 
     def exit_unary_expr(self, node: uni.UnaryExpr) -> None:
-        if node.op.name == Tok.NOT:
+        op_cls = UNARY_OP_MAP.get(node.op.name)
+        if op_cls:
             node.gen.py_ast = [
                 self.sync(
                     ast3.UnaryOp(
-                        op=self.sync(ast3.Not()),
-                        operand=cast(ast3.expr, node.operand.gen.py_ast[0]),
-                    )
-                )
-            ]
-        elif node.op.name == Tok.BW_NOT:
-            node.gen.py_ast = [
-                self.sync(
-                    ast3.UnaryOp(
-                        op=self.sync(ast3.Invert()),
-                        operand=cast(ast3.expr, node.operand.gen.py_ast[0]),
-                    )
-                )
-            ]
-        elif node.op.name == Tok.PLUS:
-            node.gen.py_ast = [
-                self.sync(
-                    ast3.UnaryOp(
-                        op=self.sync(ast3.UAdd()),
-                        operand=cast(ast3.expr, node.operand.gen.py_ast[0]),
-                    )
-                )
-            ]
-        elif node.op.name == Tok.MINUS:
-            node.gen.py_ast = [
-                self.sync(
-                    ast3.UnaryOp(
-                        op=self.sync(ast3.USub()),
+                        op=self.sync(op_cls()),
                         operand=cast(ast3.expr, node.operand.gen.py_ast[0]),
                     )
                 )
             ]
+            return
         elif node.op.name in [Tok.PIPE_FWD, Tok.KW_SPAWN, Tok.A_PIPE_FWD]:
             node.gen.py_ast = [
                 self.sync(
@@ -2573,8 +2312,10 @@ class PyastGenPass(UniPass):
 
         return JacMachineInterface.get_by_llm_call_args(self, node)
 
-    def exit_func_call(self, node: uni.FuncCall) -> None:
-        func = node.target.gen.py_ast[0]
+    def gen_call_args(
+        self, node: uni.FuncCall
+    ) -> tuple[list[ast3.expr], list[ast3.keyword]]:
+        """Generate the arguments for a function call."""
         args = []
         keywords = []
         if node.params:
@@ -2596,10 +2337,21 @@ class PyastGenPass(UniPass):
                     keywords.append(x.gen.py_ast[0])
                 else:
                     self.ice("Invalid Parameter")
-        if node.genai_call:
+        return args, keywords
+
+    def exit_func_call(self, node: uni.FuncCall) -> None:
+        if node.body_genai_call:
+            node.gen.py_ast = self.gen_llm_call_override(node)
+
+        # TODO: This needs to be changed to only generate parameters no the body.
+        elif node.genai_call:
             by_llm_call_args = self.get_by_llm_call_args(node)
             node.gen.py_ast = [self.sync(self.by_llm_call(**by_llm_call_args))]
+
         else:
+            func = node.target.gen.py_ast[0]
+            args, keywords = self.gen_call_args(node)
+
             node.gen.py_ast = [
                 self.sync(
                     ast3.Call(
@@ -2699,149 +2451,115 @@ class PyastGenPass(UniPass):
             ]
 
     def exit_edge_ref_trailer(self, node: uni.EdgeRefTrailer) -> None:
-        pynode = node.chain[0].gen.py_ast[0]
-        chomp = [*node.chain]
-        last_edge = None
+        origin = None
+        cur = node.chain[0]
+        chomp = [*node.chain[1:]]
         from_visit = bool(isinstance(node.parent, uni.VisitStmt))
-        if node.edges_only:
-            for i in node.chain:
-                if isinstance(i, uni.EdgeOpRef):
-                    last_edge = i
-        while len(chomp):
-            cur = chomp[0]
-            chomp = chomp[1:]
-            if len(chomp) == len(node.chain) - 1 and not isinstance(cur, uni.EdgeOpRef):
-                continue
-            next_i = chomp[0] if chomp else None
-            if isinstance(cur, uni.EdgeOpRef) and (
-                not next_i or not isinstance(next_i, uni.EdgeOpRef)
-            ):
-                pynode = self.translate_edge_op_ref(
-                    loc=pynode,
-                    node=cur,
-                    targ=(
-                        next_i.gen.py_ast[0]
-                        if next_i and not isinstance(next_i, uni.FilterCompr)
-                        else None
-                    ),
-                    edges_only=node.edges_only and cur == last_edge,
-                    from_visit=from_visit,
-                )
-                if next_i and isinstance(next_i, uni.FilterCompr):
-                    pynode = self.sync(
-                        ast3.Call(
-                            func=self.jaclib_obj("filter"),
-                            args=[],
-                            keywords=[
-                                self.sync(
-                                    ast3.keyword(
-                                        arg="items",
-                                        value=cast(ast3.expr, pynode),
-                                    )
-                                ),
-                                self.sync(
-                                    ast3.keyword(
-                                        arg="func",
-                                        value=cast(ast3.expr, next_i.gen.py_ast[0]),
-                                    )
-                                ),
-                            ],
-                        )
-                    )
-                chomp = chomp[1:] if next_i else chomp
-            elif isinstance(cur, uni.EdgeOpRef) and isinstance(next_i, uni.EdgeOpRef):
-                pynode = self.translate_edge_op_ref(
-                    pynode,
-                    cur,
-                    targ=None,
-                    edges_only=node.edges_only and cur == last_edge,
-                    from_visit=from_visit,
-                )
-            else:
-                raise self.ice("Invalid edge ref trailer")
-
-        node.gen.py_ast = [pynode]
-
-    def exit_edge_op_ref(self, node: uni.EdgeOpRef) -> None:
-        loc = self.sync(
-            ast3.Name(id=Con.HERE.value, ctx=ast3.Load())
-            if node.from_walker
-            else ast3.Name(id="self", ctx=ast3.Load())
-        )
-        node.gen.py_ast = [loc]
 
-    def translate_edge_op_ref(
-        self,
-        loc: ast3.AST,
-        node: uni.EdgeOpRef,
-        targ: ast3.AST | None,
-        edges_only: bool,
-        from_visit: bool,
-    ) -> ast3.AST:
-        """Generate ast for edge op ref call."""
-        keywords = [self.sync(ast3.keyword(arg="sources", value=cast(ast3.expr, loc)))]
+        if not isinstance(cur, uni.EdgeOpRef):
+            origin = cur.gen.py_ast[0]
+            cur = cast(uni.EdgeOpRef, chomp.pop(0))
 
-        if targ:
-            keywords.append(
-                self.sync(ast3.keyword(arg="targets", value=cast(ast3.expr, targ)))
+        pynode = self.sync(
+            ast3.Call(
+                func=self.jaclib_obj("Path"),
+                args=[cast(ast3.expr, origin or cur.gen.py_ast[0])],
+                keywords=[],
             )
+        )
 
-        if node.edge_dir != EdgeDir.OUT:
-            keywords.append(
-                self.sync(
-                    ast3.keyword(
-                        arg="dir",
-                        value=self.sync(
-                            ast3.Attribute(
-                                value=self.jaclib_obj("EdgeDir"),
-                                attr=node.edge_dir.name,
-                                ctx=ast3.Load(),
-                            )
-                        ),
+        while True:
+            keywords = []
+            if cur.filter_cond:
+                keywords.append(
+                    self.sync(
+                        ast3.keyword(
+                            arg="edge",
+                            value=cast(
+                                ast3.expr, self.sync(cur.filter_cond.gen.py_ast[0])
+                            ),
+                        )
                     )
                 )
-            )
 
-        if node.filter_cond:
-            keywords.append(
-                self.sync(
-                    ast3.keyword(
-                        arg="filter",
-                        value=cast(
-                            ast3.expr, self.sync(node.filter_cond.gen.py_ast[0])
-                        ),
+            if chomp and not isinstance(chomp[0], uni.EdgeOpRef):
+                filt = chomp.pop(0)
+                keywords.append(
+                    self.sync(
+                        ast3.keyword(
+                            arg="node",
+                            value=cast(ast3.expr, self.sync(filt.gen.py_ast[0])),
+                        )
                     )
                 )
+
+            pynode = self.sync(
+                ast3.Call(
+                    func=self.sync(
+                        ast3.Attribute(
+                            value=pynode,
+                            attr=f"_{cur.edge_dir.name.lower()}",
+                            ctx=ast3.Load(),
+                        )
+                    ),
+                    args=[],
+                    keywords=keywords,
+                )
             )
 
-        if edges_only:
-            keywords.append(
-                self.sync(
-                    ast3.keyword(
-                        arg="edges_only",
-                        value=self.sync(ast3.Constant(value=edges_only)),
-                    )
+            if chomp:
+                cur = cast(uni.EdgeOpRef, chomp.pop(0))
+            else:
+                break
+
+        if node.edges_only:
+            pynode = self.sync(
+                ast3.Call(
+                    func=self.sync(
+                        ast3.Attribute(
+                            value=pynode,
+                            attr="edge",
+                            ctx=ast3.Load(),
+                        )
+                    ),
+                    args=[],
+                    keywords=[],
                 )
             )
 
         if from_visit:
-            keywords.append(
-                self.sync(
-                    ast3.keyword(
-                        arg="from_visit",
-                        value=self.sync(ast3.Constant(value=from_visit)),
-                    )
+            pynode = self.sync(
+                ast3.Call(
+                    func=self.sync(
+                        ast3.Attribute(
+                            value=pynode,
+                            attr="visit",
+                            ctx=ast3.Load(),
+                        )
+                    ),
+                    args=[],
+                    keywords=[],
                 )
             )
 
-        return self.sync(
+        pynode = self.sync(
             ast3.Call(
                 func=self.jaclib_obj("refs"),
-                args=[],
-                keywords=keywords,
+                args=[pynode],
+                keywords=[],
             )
         )
 
+        node.gen.py_ast = [pynode]
+
+    def exit_edge_op_ref(self, node: uni.EdgeOpRef) -> None:
+        loc = self.sync(
+            ast3.Name(id=Con.HERE.value, ctx=ast3.Load())
+            if node.from_walker
+            else ast3.Name(id="self", ctx=ast3.Load())
+        )
+        node.gen.py_ast = [loc]
+
     def exit_disconnect_op(self, node: uni.DisconnectOp) -> None:
         node.gen.py_ast = node.edge_spec.gen.py_ast
 
@@ -3107,62 +2825,9 @@ class PyastGenPass(UniPass):
         ]
 
     def exit_token(self, node: uni.Token) -> None:
-        if node.name == Tok.KW_AND:
-            node.gen.py_ast = [self.sync(ast3.And())]
-        elif node.name == Tok.KW_OR:
-            node.gen.py_ast = [self.sync(ast3.Or())]
-        elif node.name in [Tok.PLUS, Tok.ADD_EQ]:
-            node.gen.py_ast = [self.sync(ast3.Add())]
-        elif node.name in [Tok.BW_AND, Tok.BW_AND_EQ]:
-            node.gen.py_ast = [self.sync(ast3.BitAnd())]
-        elif node.name in [Tok.BW_OR, Tok.BW_OR_EQ]:
-            node.gen.py_ast = [self.sync(ast3.BitOr())]
-        elif node.name in [Tok.BW_XOR, Tok.BW_XOR_EQ]:
-            node.gen.py_ast = [self.sync(ast3.BitXor())]
-        elif node.name in [Tok.DIV, Tok.DIV_EQ]:
-            node.gen.py_ast = [self.sync(ast3.Div())]
-        elif node.name in [Tok.FLOOR_DIV, Tok.FLOOR_DIV_EQ]:
-            node.gen.py_ast = [self.sync(ast3.FloorDiv())]
-        elif node.name in [Tok.LSHIFT, Tok.LSHIFT_EQ]:
-            node.gen.py_ast = [self.sync(ast3.LShift())]
-        elif node.name in [Tok.MOD, Tok.MOD_EQ]:
-            node.gen.py_ast = [self.sync(ast3.Mod())]
-        elif node.name in [Tok.STAR_MUL, Tok.MUL_EQ]:
-            node.gen.py_ast = [self.sync(ast3.Mult())]
-        elif node.name in [Tok.DECOR_OP, Tok.MATMUL_EQ]:
-            node.gen.py_ast = [self.sync(ast3.MatMult())]
-        elif node.name in [Tok.STAR_POW, Tok.STAR_POW_EQ]:
-            node.gen.py_ast = [self.sync(ast3.Pow())]
-        elif node.name in [Tok.RSHIFT, Tok.RSHIFT_EQ]:
-            node.gen.py_ast = [self.sync(ast3.RShift())]
-        elif node.name in [Tok.MINUS, Tok.SUB_EQ]:
-            node.gen.py_ast = [self.sync(ast3.Sub())]
-        elif node.name in [Tok.BW_NOT, Tok.BW_NOT_EQ]:
-            node.gen.py_ast = [self.sync(ast3.Invert())]
-        elif node.name in [Tok.NOT]:
-            node.gen.py_ast = [self.sync(ast3.Not())]
-        elif node.name in [Tok.EQ]:
-            node.gen.py_ast = [self.sync(ast3.NotEq())]
-        elif node.name == Tok.EE:
-            node.gen.py_ast = [self.sync(ast3.Eq())]
-        elif node.name == Tok.GT:
-            node.gen.py_ast = [self.sync(ast3.Gt())]
-        elif node.name == Tok.GTE:
-            node.gen.py_ast = [self.sync(ast3.GtE())]
-        elif node.name == Tok.KW_IN:
-            node.gen.py_ast = [self.sync(ast3.In())]
-        elif node.name == Tok.KW_IS:
-            node.gen.py_ast = [self.sync(ast3.Is())]
-        elif node.name == Tok.KW_ISN:
-            node.gen.py_ast = [self.sync(ast3.IsNot())]
-        elif node.name == Tok.LT:
-            node.gen.py_ast = [self.sync(ast3.Lt())]
-        elif node.name == Tok.LTE:
-            node.gen.py_ast = [self.sync(ast3.LtE())]
-        elif node.name == Tok.NE:
-            node.gen.py_ast = [self.sync(ast3.NotEq())]
-        elif node.name == Tok.KW_NIN:
-            node.gen.py_ast = [self.sync(ast3.NotIn())]
+        op_cls = TOKEN_AST_MAP.get(node.name)
+        if op_cls:
+            node.gen.py_ast = [self.sync(op_cls())]
 
     def exit_name(self, node: uni.Name) -> None:
         node.gen.py_ast = [
@@ -3173,21 +2838,7 @@ class PyastGenPass(UniPass):
         node.gen.py_ast = [self.sync(ast3.Constant(value=float(node.value)))]
 
     def exit_int(self, node: uni.Int) -> None:
-        def handle_node_value(value: str) -> int:
-            if value.startswith(("0x", "0X")):
-                return int(value, 16)
-            elif value.startswith(("0b", "0B")):
-                return int(value, 2)
-            elif value.startswith(("0o", "0O")):
-                return int(value, 8)
-            else:
-                return int(value)
-
-        node.gen.py_ast = [
-            self.sync(
-                ast3.Constant(value=handle_node_value(str(node.value)), kind=None)
-            )
-        ]
+        node.gen.py_ast = [self.sync(ast3.Constant(value=int(node.value, 0)))]
 
     def exit_string(self, node: uni.String) -> None:
         node.gen.py_ast = [self.sync(ast3.Constant(value=node.lit_value))]