libcst: Inserting new node adds inline code and a semicolon

Question

I am trying to introduce a new node (as a new line of code), exactly before an Assign node.

The issue occurs when using FlattenSentinel to introduce the new node as I want the node to be separate, but libcst concatenates them using a semicolon (;), example:

a = 6

Becomes:

print('returning'); a = 6

Code to reproduce example:

import libcst as cst
class MyTransformer(cst.CSTTransformer):

    def leave_Assign(self, old_node, updated_node):
        log_stmt = cst.Expr(cst.parse_expression("print('returning')"))
        return cst.FlattenSentinel([log_stmt, updated_node])

source_tree = cst.parse_module("a = 6")
modified_tree = source_tree.visit(MyTransformer())
print(modified_tree.code)

I also tried introducing a new line but is looks even worse, code sample:

def leave_Assign(self, old_node, updated_node):
    log_stmt = cst.Expr(cst.parse_expression("print('returning')"))
    return cst.FlattenSentinel([log_stmt, cst.Expr(cst.Newline()), updated_node])

My desired result would be to insert the new node above the existing node (at same level ), without the semicolon, like this:

print('returning')
a = 6

Is this possible in libcst?

Answer 1

Your AST is like this:

Module
  FunctionDef
    SimpleStatementLine
      Assign
        Target
        Integer

When you try to replace the Assign with multiple nodes, you still need to respect the current line, which forces the semicolons.

Module
  FunctionDef
    SimpleStatementLine
      Call // Your print call
      Assign // a = 6
        Target
        Integer

Instead, you want to replace the parent SimpleStatementLine with a FlattenSentinel of new SimpleStatementLines with your new nodes. So, you actually have to modify your code to operate on leave_SimpleStatementLine.

Answer 2

Ampersandy provides explanation. You can search for SimpleStatementLine nodes with a Assign child and insert your nodes before founded node with FlattenSentinel. Here the code:

import libcst as cst


class MyTransformer(cst.CSTTransformer):
    def leave_SimpleStatementLine(self, old_node, updated_node):
        if any(isinstance(child, cst.Assign) for child in updated_node.children):
            log_stmt = cst.parse_module("print('returning')").body
            return cst.FlattenSentinel([*log_stmt, updated_node])
        return updated_node


source_code = """
def func():
    a = 6
a = 2
"""

source_tree = cst.parse_module(source_code)
modified_tree = source_tree.visit(MyTransformer())
print(modified_tree.code)

which produces:

def func():
    print('returning')
    a = 6
print('returning')
a = 2