jaclang 0.8.6__py3-none-any.whl → 0.8.8__py3-none-any.whl

jaclang/compiler/parser.py

@@ -5,12 +5,14 @@ from __future__ import annotations
 import keyword
 import logging
 import os
+import sys
 from typing import Callable, Sequence, TYPE_CHECKING, TypeAlias, TypeVar, cast
 
 import jaclang.compiler.unitree as uni
-from jaclang.compiler import jac_lark as jl
+from jaclang.compiler import TOKEN_MAP, jac_lark as jl
 from jaclang.compiler.constant import EdgeDir, Tokens as Tok
 from jaclang.compiler.passes.main import Transform
+from jaclang.utils.helpers import ANSIColors
 from jaclang.vendor.lark import Lark, Transformer, Tree, logger
 
 if TYPE_CHECKING:
@@ -39,22 +41,17 @@ class JacParser(Transform[uni.Source, uni.Module]):
             tree, comments = JacParser.parse(ir_in.value, on_error=self.error_callback)
             mod = JacParser.TreeToAST(parser=self).transform(tree)
             ir_in.comments = [self.proc_comment(i, mod) for i in comments]
-            if isinstance(mod, uni.Module):
-                self.ir_out = mod
-                return mod
-            else:
+            if not isinstance(mod, uni.Module):
                 raise self.ice()
+            if len(self.errors_had) != 0:
+                mod.has_syntax_errors = True
+                self.report_errors()
+            self.ir_out = mod
+            return mod
         except jl.UnexpectedInput as e:
-            catch_error = uni.EmptyToken()
-            catch_error.orig_src = ir_in
-            catch_error.line_no = e.line
-            catch_error.end_line = e.line
-            catch_error.c_start = e.column
-            catch_error.c_end = e.column + 1
-            catch_error.pos_start = e.pos_in_stream or 0
-            catch_error.pos_end = catch_error.pos_start + 1
+            catch_error = self.error_to_token(e)
+            error_msg = self.error_to_message(e)
 
-            error_msg = "Syntax Error"
             if len(e.args) >= 1 and isinstance(e.args[0], str):
                 error_msg += e.args[0]
             self.log_error(error_msg, node_override=catch_error)
@@ -62,7 +59,12 @@ class JacParser(Transform[uni.Source, uni.Module]):
         except Exception as e:
             raise e
 
-        return uni.Module.make_stub(inject_src=ir_in)
+        # If we reach here, there was a syntax error, mark the module as such
+        # and report errors.
+        self.report_errors()
+        mod = uni.Module.make_stub(inject_src=ir_in)
+        mod.has_syntax_errors = True
+        return mod
 
     @staticmethod
     def proc_comment(token: jl.Token, mod: uni.UniNode) -> uni.CommentToken:
@@ -80,10 +82,109 @@ class JacParser(Transform[uni.Source, uni.Module]):
             kid=[],
         )
 
+    _MISSING_TOKENS = [
+        Tok.SEMI,
+        Tok.COMMA,
+        Tok.COLON,
+        Tok.RPAREN,
+        Tok.RBRACE,
+        Tok.RSQUARE,
+        Tok.RETURN_HINT,
+    ]
+
     def error_callback(self, e: jl.UnexpectedInput) -> bool:
         """Handle error."""
+        iparser = e.interactive_parser
+
+        def try_feed_missing_token(iparser: jl.InteractiveParser) -> Tok | None:
+            """Feed a missing token to the parser."""
+            # If any of the below token is missing, insert them and continue parsing.
+            accepts = iparser.accepts()
+            for tok in JacParser._MISSING_TOKENS:
+                if tok.name in accepts:
+                    iparser.feed_token(jl.Token(tok.name, TOKEN_MAP[tok.name]))
+                    return tok
+            return None
+
+        def feed_current_token(iparser: jl.InteractiveParser, tok: jl.Token) -> bool:
+            """Feed the current token to the parser."""
+            while tok.type not in iparser.accepts():
+                if not try_feed_missing_token(iparser):
+                    return False
+            iparser.feed_token(tok)
+            return True
+
+        if isinstance(e, jl.UnexpectedToken):
+            # If last token is DOT and we expect a NAME, insert a NAME token
+            last_tok: jl.Token | None = (
+                e.token_history[-1]
+                if e.token_history and len(e.token_history) >= 1
+                else None
+            )
+            if (
+                last_tok
+                and last_tok.type == Tok.DOT.name
+                and (Tok.NAME.name in e.accepts)
+            ):
+                self.log_error("Incomplete member access", self.error_to_token(e))
+                iparser.feed_token(jl.Token(Tok.NAME.name, "recover_name_token"))
+                return feed_current_token(iparser, e.token)
+
+            # We're calling try_feed_missing_token twice here because the first missing
+            # will be reported as such and we don't for the consequent missing token.
+            if tk := try_feed_missing_token(iparser):
+                self.log_error(f"Missing {tk.name}", self.error_to_token(e))
+                return feed_current_token(iparser, e.token)
+
+            # Ignore unexpected tokens and continue parsing till we reach a known state.
+            self.log_error(
+                f"Unexpected token '{e.token.value}'", self.error_to_token(e)
+            )
+            return True
+
         return False
 
+    def error_to_message(self, e: jl.UnexpectedInput) -> str:
+        """Return an error message based on the exception."""
+        # TODO: Match more specific errors with lark's example based matching.
+        # Reference: https://github.com/lark-parser/lark/blob/master/examples/advanced/error_reporting_lalr.py
+        # e.match_examples()
+        if isinstance(e, jl.UnexpectedToken):
+            return f"Unexpected token '{e.token.value}'"
+        return "Syntax Error"
+
+    def error_to_token(self, e: jl.UnexpectedInput) -> uni.Token:
+        """Convert error to token."""
+        catch_error = uni.EmptyToken()
+        catch_error.orig_src = self.ir_in
+        catch_error.line_no = e.line
+        catch_error.end_line = e.line
+        catch_error.c_start = e.column
+        catch_error.pos_start = e.pos_in_stream or 0
+        if isinstance(e, jl.UnexpectedToken) and e.token:
+            catch_error.c_end = e.token.end_column or (e.column + 1)
+            catch_error.pos_end = e.token.end_pos or (catch_error.pos_start + 1)
+        else:
+            catch_error.c_end = e.column + 1
+            catch_error.pos_end = catch_error.pos_start + 1
+        return catch_error
+
+    def report_errors(self, *, colors: bool = True) -> None:
+        """Report errors to the user."""
+        # TODO: Write a better IO system.
+        # NOTE: Currently it writes all the errors to stderr cause LSP JsonRPC uses stdout for IPC.
+        if not sys.stderr.isatty():
+            # FIXME: If we're outputting to a file (pipe, redirection, etc) other
+            # than a terminal we disable colors however we should be able to force
+            # colors with a configuration.
+            colors = False
+        for alrt in self.errors_had:
+            error_label = (
+                "Error:" if not colors else f"{ANSIColors.RED}Error:{ANSIColors.END}"
+            )
+            print(error_label, end=" ", file=sys.stderr)
+            print(alrt.pretty_print(colors=colors), file=sys.stderr)
+
     @staticmethod
     def _comment_callback(comment: jl.Token) -> None:
         JacParser.comment_cache.append(comment)
@@ -848,25 +949,104 @@ class JacParser(Transform[uni.Source, uni.Module]):
             if self.match_token(Tok.RETURN_HINT):
                 return_spec = self.consume(uni.Expr)
                 return uni.FuncSignature(
+                    posonly_params=[],
                     params=[],
+                    varargs=None,
+                    kwonlyargs=[],
+                    kwargs=None,
                     return_type=return_spec,
                     kid=self.flat_cur_nodes,
                 )
             # Otherwise, parse the traditional parameter list form
             else:
                 self.consume_token(Tok.LPAREN)
-                params = self.match(list)
+                all_params = self.match(list) or []
+                posonly_params, params, varargs, kwonlyargs, kwargs = (
+                    self._parse_parameter_categories(all_params)
+                )
                 self.consume_token(Tok.RPAREN)
                 if self.match_token(Tok.RETURN_HINT):
                     return_spec = self.consume(uni.Expr)
                 return uni.FuncSignature(
-                    params=(
-                        self.extract_from_list(params, uni.ParamVar) if params else []
-                    ),
+                    posonly_params=posonly_params,
+                    params=params,
+                    varargs=varargs,
+                    kwonlyargs=kwonlyargs,
+                    kwargs=kwargs,
                     return_type=return_spec,
                     kid=self.flat_cur_nodes,
                 )
 
+        def _parse_parameter_categories(self, all_params: list[uni.UniNode]) -> tuple[
+            list[uni.ParamVar],
+            list[uni.ParamVar],
+            uni.ParamVar | None,
+            list[uni.ParamVar],
+            uni.ParamVar | None,
+        ]:
+            posonly_params = []
+            params = []
+            varargs = None
+            kwonlyargs = []
+            kwargs = None
+
+            # Initial state determination
+            cur_state = "positional"
+            for param in all_params:
+                if isinstance(param, uni.Token) and param.name == Tok.DIV:
+                    cur_state = "posonly"
+                    break
+
+            for cur_nd in all_params:
+                cur_state = self._update_parameter_state(cur_nd, cur_state)
+                if isinstance(cur_nd, uni.ParamVar):
+                    if cur_state == "positional":
+                        cur_nd.param_kind = uni.ParamKind.NORMAL
+                        params.append(cur_nd)
+                    elif cur_state == "posonly":
+                        cur_nd.param_kind = uni.ParamKind.POSONLY
+                        posonly_params.append(cur_nd)
+                    elif cur_state == "varargs":
+                        cur_nd.param_kind = uni.ParamKind.VARARG
+                        varargs = cur_nd
+                        cur_state = "keyword_only"
+                    elif cur_state == "keyword_only":
+                        cur_nd.param_kind = uni.ParamKind.KWONLY
+                        kwonlyargs.append(cur_nd)
+                    elif cur_state == "kwargs":
+                        cur_nd.param_kind = uni.ParamKind.KWARG
+                        kwargs = cur_nd
+                    else:
+                        raise self.ice()
+
+            return posonly_params, params, varargs, kwonlyargs, kwargs
+
+        def _update_parameter_state(self, cur_nd: uni.UniNode, cur_state: str) -> str:
+            if isinstance(cur_nd, uni.Token):
+                if cur_nd.name == Tok.DIV:
+                    if cur_state in ["keyword_only", "kwargs", "positional"]:
+                        self.parse_ref.log_error(
+                            "Invalid syntax in function parameters: '/' cannot appear after '*' or '**'.",
+                            node_override=cur_nd,
+                        )
+                    return "positional"
+                elif cur_nd.name == Tok.STAR_MUL:
+                    if cur_state in ["keyword_only", "kwargs"]:
+                        self.parse_ref.log_error(
+                            "Invalid syntax in function parameters: '*' cannot appear after '**'.",
+                            node_override=cur_nd,
+                        )
+                    return "keyword_only"
+                elif cur_nd.name == Tok.COMMA:
+                    return cur_state
+
+            elif isinstance(cur_nd, uni.ParamVar):
+                if cur_nd.is_vararg:
+                    return "varargs"
+                if cur_nd.is_kwargs:
+                    return "kwargs"
+            return cur_state
+
         def func_decl_params(self, _: None) -> list[uni.UniNode]:
             """Grammar rule.
 
@@ -874,11 +1054,17 @@ class JacParser(Transform[uni.Source, uni.Module]):
             """
             return self.cur_nodes
 
-        def param_var(self, _: None) -> uni.ParamVar:
+        def param_var(self, _: None) -> uni.ParamVar | uni.Token:
             """Grammar rule.
 
-            param_var: (STAR_POW | STAR_MUL)? NAME type_tag (EQ expression)?
+            param_var: (STAR_POW | STAR_MUL)? named_ref type_tag (EQ expression)?
+                    | DIV
+                    | STAR_MUL
             """
+            if len(self.cur_nodes) == 1 and (
+                star_only := self.match_token(Tok.DIV) or self.match_token(Tok.STAR_MUL)
+            ):
+                return star_only
             star = self.match_token(Tok.STAR_POW) or self.match_token(Tok.STAR_MUL)
             name = self.consume(uni.Name)
             type_tag = self.consume(uni.SubTag)
@@ -1028,7 +1214,6 @@ class JacParser(Transform[uni.Source, uni.Module]):
                     | (yield_expr | KW_YIELD) SEMI
                     | raise_stmt SEMI
                     | assert_stmt SEMI
-                    | check_stmt SEMI
                     | assignment SEMI
                     | delete_stmt SEMI
                     | report_stmt SEMI
@@ -1307,18 +1492,6 @@ class JacParser(Transform[uni.Source, uni.Module]):
                 kid=self.cur_nodes,
             )
 
-        def check_stmt(self, _: None) -> uni.CheckStmt:
-            """Grammar rule.
-
-            check_stmt: KW_CHECK expression
-            """
-            self.consume_token(Tok.KW_CHECK)
-            target = self.consume(uni.Expr)
-            return uni.CheckStmt(
-                target=target,
-                kid=self.cur_nodes,
-            )
-
         def ctrl_stmt(self, _: None) -> uni.CtrlStmt | uni.DisengageStmt:
             """Grammar rule.
 
@@ -1545,9 +1718,13 @@ class JacParser(Transform[uni.Source, uni.Module]):
                 sig_kid.append(return_type)
             signature = (
                 uni.FuncSignature(
+                    posonly_params=[],
                     params=(
                         self.extract_from_list(params, uni.ParamVar) if params else []
                     ),
+                    varargs=None,
+                    kwonlyargs=[],
+                    kwargs=None,
                     return_type=return_type,
                     kid=sig_kid,
                 )
@@ -2328,6 +2505,7 @@ class JacParser(Transform[uni.Source, uni.Module]):
             LSQUARE (KW_NODE| KW_EDGE)? expression? (edge_op_ref (filter_compr | expression)?)+ RSQUARE
             """
             self.consume_token(Tok.LSQUARE)
+            is_async = bool(self.match_token(Tok.KW_ASYNC))
             edges_only = bool(self.match_token(Tok.KW_EDGE))
             self.match_token(Tok.KW_NODE)
             valid_chain = []
@@ -2338,6 +2516,7 @@ class JacParser(Transform[uni.Source, uni.Module]):
             return uni.EdgeRefTrailer(
                 chain=valid_chain,
                 edges_only=edges_only,
+                is_async=is_async,
                 kid=self.cur_nodes,
             )
 

jaclang/compiler/passes/main/__init__.py

@@ -4,7 +4,6 @@ from ..transform import Alert, Transform  # noqa: I100
 from .annex_pass import JacAnnexPass  # noqa: I100
 from .binder_pass import BinderPass  # noqa: I100
 from .sym_tab_build_pass import SymTabBuildPass, UniPass  # noqa: I100
-from .sym_tab_link_pass import SymTabLinkPass  # noqa: I100
 from .def_use_pass import DefUsePass  # noqa: I100
 from .sem_def_match_pass import SemDefMatchPass  # noqa: I100
 from .import_pass import JacImportDepsPass  # noqa: I100
@@ -12,10 +11,10 @@ from .def_impl_match_pass import DeclImplMatchPass  # noqa: I100
 from .type_checker_pass import TypeCheckPass  # noqa: I100
 from .pyast_load_pass import PyastBuildPass  # type: ignore # noqa: I100
 from .pyast_gen_pass import PyastGenPass  # noqa: I100
+from .predynamo_pass import PreDynamoPass  # noqa: I100
 from .pybc_gen_pass import PyBytecodeGenPass  # noqa: I100
 from .cfg_build_pass import CFGBuildPass  # noqa: I100
 from .pyjac_ast_link_pass import PyJacAstLinkPass  # noqa: I100
-from .inheritance_pass import InheritancePass  # noqa: I100
 
 
 __all__ = [
@@ -28,14 +27,13 @@ __all__ = [
     "BinderPass",
     "TypeCheckPass",
     "SymTabBuildPass",
-    "SymTabLinkPass",
     "DeclImplMatchPass",
     "DefUsePass",
     "SemDefMatchPass",
     "PyastBuildPass",
     "PyastGenPass",
+    "PreDynamoPass",
     "PyBytecodeGenPass",
     "CFGBuildPass",
     "PyJacAstLinkPass",
-    "InheritancePass",
 ]

jaclang/compiler/passes/main/def_use_pass.py

@@ -30,7 +30,6 @@ class DefUsePass(UniPass):
     """Jac Ast build pass."""
 
     def enter_archetype(self, node: uni.Archetype) -> None:
-        node.sym_tab.inherit_baseclasses_sym(node)
 
         def inform_from_walker(node: uni.UniNode) -> None:
             for i in (
@@ -47,9 +46,6 @@ class DefUsePass(UniPass):
                 if isinstance(i.body, uni.ImplDef):
                     inform_from_walker(i.body)
 
-    def enter_enum(self, node: uni.Enum) -> None:
-        node.sym_tab.inherit_baseclasses_sym(node)
-
     def enter_type_ref(self, node: uni.TypeRef) -> None:
         node.sym_tab.use_lookup(node)
 

jaclang/compiler/passes/main/predynamo_pass.py

@@ -0,0 +1,221 @@
+"""Pytorch Fix Pass."""
+
+import ast as ast3
+from typing import Optional, TypeVar, cast
+
+import jaclang.compiler.unitree as uni
+from jaclang.compiler.constant import Tokens as Tok
+from jaclang.compiler.passes import UniPass
+
+
+T = TypeVar("T", bound=ast3.AST)
+
+
+class PreDynamoPass(UniPass):
+    """Pre-Dynamo pass for PyTorch."""
+
+    def enter_node(self, node: uni.UniNode) -> None:
+        """Enter node."""
+        super().enter_node(node)
+
+    def exit_node(self, node: uni.UniNode) -> None:
+        """Exit node."""
+        super().exit_node(node)
+
+    def gen_name(self, node: uni.UniNode, name: Tok, value: str) -> uni.Name:
+        """Generate Name."""
+        return uni.Name(
+            name=name,
+            value=value,
+            orig_src=node.loc.orig_src,
+            col_start=node.loc.col_start,
+            col_end=0,
+            line=node.loc.first_line,
+            end_line=node.loc.last_line,
+            pos_start=0,
+            pos_end=0,
+        )
+
+    def replace_node(
+        self,
+        new_nodes: list[uni.UniNode] | uni.UniNode,
+        old_node: uni.UniNode,
+        attr: str,
+    ) -> None:
+        """Replace old node with new nodes in parent's body and kid lists."""
+        parent = old_node.parent
+        if isinstance(new_nodes, uni.UniNode):
+            new_nodes.parent = parent
+            if hasattr(parent, attr):
+                lst = getattr(parent, attr)
+                if old_node in lst:
+                    idx = lst.index(old_node)
+                    lst[idx] = new_nodes
+            if hasattr(parent, "kid") and old_node in parent.kid:
+                idx = parent.kid.index(old_node)
+                parent.kid[idx] = new_nodes
+        else:  # list of nodes
+            for n in new_nodes:
+                n.parent = parent
+            if hasattr(parent, attr):
+                lst = getattr(parent, attr)
+                if old_node in lst:
+                    idx = lst.index(old_node)
+                    setattr(parent, attr, lst[:idx] + new_nodes + lst[idx + 1 :])
+            if hasattr(parent, "kid") and old_node in parent.kid:
+                idx = parent.kid.index(old_node)
+                parent.kid = parent.kid[:idx] + new_nodes + parent.kid[idx + 1 :]
+
+    def check_same_lhs(
+        self, assign_a: uni.UniNode, assign_b: uni.UniNode
+    ) -> Optional[uni.Name]:
+        """Return the common LHS target if both are simple assignment with same target."""
+        if not (
+            isinstance(assign_a, uni.Assignment)
+            and isinstance(assign_b, uni.Assignment)
+        ):
+            return None
+        ta, tb = assign_a.target[0], assign_b.target[0]
+        if not (isinstance(ta, uni.Name) and isinstance(tb, uni.Name)):
+            return None
+        if ta.value != tb.value:
+            return None
+        return ta  # common target
+
+    def check_call(self, node: uni.ExprStmt) -> Optional[tuple]:
+        """Return (target, name, tensor_expr, kwargs) if node is target(name, tensor_expr, **kwargs)."""
+        if isinstance(node, uni.ExprStmt) and isinstance(node.expr, uni.FuncCall):
+            call = node.expr
+            if (
+                isinstance(call.target, uni.AtomTrailer)
+                and len(call.params) >= 2
+                and isinstance(call.params[0], (uni.String, uni.MultiString))
+                and isinstance(call.params[1], uni.Expr)
+            ):
+                name = (
+                    call.params[0]
+                    if isinstance(call.params[0], uni.String)
+                    else call.params[0].strings[0]
+                )
+                tensor_expr = call.params[1]
+                kwargs = (
+                    {
+                        kw.key._sym_name: kw.value
+                        for kw in call.params[2:]
+                        if isinstance(kw, uni.KWPair)
+                    }
+                    if len(call.params) > 2
+                    else {}
+                )
+                return (call.target, name, tensor_expr, kwargs)
+        return None
+
+    def exit_if_stmt(self, node: uni.IfStmt) -> None:
+        """Exit if statement."""
+        a0 = node.body[0]
+        new_node = None
+        if node.else_body:
+            b0 = node.else_body.body[0]
+        else:
+            return
+        if isinstance(a0, uni.Assignment) and isinstance(b0, uni.Assignment):
+            lhs = self.check_same_lhs(a0, b0)
+            if lhs is not None:
+                func_name = self.gen_name(node, Tok.NAME, "torch")
+                attr_name = self.gen_name(node, Tok.NAME, "where")
+                target = uni.AtomTrailer(
+                    target=func_name,
+                    right=attr_name,
+                    is_attr=True,
+                    is_null_ok=False,
+                    kid=[func_name, attr_name],
+                )
+                call = uni.FuncCall(
+                    target=target,
+                    params=[
+                        node.condition,
+                        cast(uni.Expr, a0.value),
+                        cast(uni.Expr, b0.value),
+                    ],
+                    genai_call=None,
+                    kid=[target, node.condition, a0, b0],
+                )
+                new_node = uni.Assignment(
+                    target=[lhs], value=call, type_tag=None, kid=[lhs, call]
+                )
+                self.replace_node(new_node, node, "body")
+
+        elif isinstance(a0, uni.ReturnStmt) and isinstance(b0, uni.ReturnStmt):
+            aexpr, bexpr = a0.expr, b0.expr
+            if aexpr is None or bexpr is None:
+                return
+            func_name = self.gen_name(node, Tok.NAME, "torch")
+            attr_name = self.gen_name(node, Tok.NAME, "where")
+            target = uni.AtomTrailer(
+                target=func_name,
+                right=attr_name,
+                is_attr=True,
+                is_null_ok=False,
+                kid=[func_name, attr_name],
+            )
+            call = uni.FuncCall(
+                target=target,
+                params=[node.condition, cast(uni.Expr, aexpr), cast(uni.Expr, bexpr)],
+                genai_call=None,
+                kid=[target, node.condition, a0, b0],
+            )
+            new_node = uni.ReturnStmt(expr=call, kid=[call])
+            self.replace_node(new_node, node, "body")
+
+        elif isinstance(a0, uni.ExprStmt) and isinstance(b0, uni.ExprStmt):
+            a_reg = self.check_call(a0)
+            b_reg = self.check_call(b0)
+            if a_reg is not None and b_reg is not None:
+                a_target, a_name, a_expr, a_kwargs = a_reg
+                b_target, b_name, b_expr, b_kwargs = b_reg
+                if a_name.value == b_name.value and set(a_kwargs.keys()) == set(
+                    b_kwargs.keys()
+                ):
+                    tmp_name = self.gen_name(node, Tok.NAME, f"__{eval(a_name.value)}")
+                    tmp_name.py_ctx_func = ast3.Store
+                    func_name = self.gen_name(node, Tok.NAME, "torch")
+                    attr_name = self.gen_name(node, Tok.NAME, "where")
+                    target = uni.AtomTrailer(
+                        target=func_name,
+                        right=attr_name,
+                        is_attr=True,
+                        is_null_ok=False,
+                        kid=[func_name, attr_name],
+                    )
+                    call = uni.FuncCall(
+                        target=target,
+                        params=[node.condition, a_expr, b_expr],
+                        genai_call=None,
+                        kid=[target, node.condition, a_expr, b_expr],
+                    )
+                    assign_node = uni.Assignment(
+                        target=[tmp_name],
+                        value=call,
+                        type_tag=None,
+                        kid=[tmp_name, call],
+                    )
+
+                    kwargs_nodes = [
+                        uni.KWPair(
+                            name := self.gen_name(node, Tok.NAME, k), v, [name, v]
+                        )
+                        for k, v in a_kwargs.items()
+                    ]
+                    param_name = self.gen_name(
+                        node, Tok.NAME, f"__{eval(a_name.value)}"
+                    )
+                    reg_call = uni.FuncCall(
+                        target=a_target,
+                        params=[a_name, param_name] + kwargs_nodes,
+                        genai_call=None,
+                        kid=[a_target, a_name, param_name] + kwargs_nodes,
+                    )
+                    reg_node = uni.ExprStmt(
+                        expr=reg_call, in_fstring=False, kid=[reg_call]
+                    )
+                    self.replace_node([assign_node, reg_node], node, "body")