libcst: при вставке нового узла добавляется встроенный код и точка с запятой

Question

Я пытаюсь ввести новый узел (как новую строку кода) непосредственно перед узлом Assign.

Проблема возникает при использовании FlattenSentinel чтобы представить новый узел, поскольку я хочу, чтобы узел был отдельным, но libcst объединяет их, используя точку с запятой ( ;), пример:

      a = 6

Становится:

      print('returning'); a = 6

Код для выше:

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

Я также пробовал ввести новую строку, но выглядит еще хуже, пример кода:

      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])

Моим желаемым результатом было бы вставить новый узел над существующим узлом (на том же уровне) без точки с запятой, например:

      print('returning')
a = 6

Возможно ли это в libcst?

Answer 1

Ampersandy дает объяснение. Вы можете искатьSimpleStatementLineузлы сAssignдочерний элемент и вставьте свои узлы перед найденным узлом с помощьюFlattenSentinel. Вот код:

      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)

который производит:

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

Answer 2

Ваш AST выглядит следующим образом:

      Module
  FunctionDef
    SimpleStatementLine
      Assign
        Target
        Integer

Когда вы пытаетесь заменить Assign несколькими узлами, вам все равно нужно соблюдать текущую строку, что заставляет использовать точки с запятой.

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

Вместо этого вы хотите заменить родителя на FlattenSentinelиз нового SimpleStatementLines с вашими новыми узлами. Итак, вам действительно нужно изменить свой код, чтобы работать с leave_SimpleStatementLine.