is there a better way to do this with libcst rn

[dheerajck]

Input

# 1
[
    1,  # comment 1
    2,  # comment 2
    3,  # comment 3
]

# 2
[1, 2, 3]

Output

# 1
[
  1, # comment 1
  2, # comment 2
  3, # comment 3
  4, # comment 4
]

# 2
[1, 2, 3, 4,  # comment 4]

import libcst as cst


class ListAppender(cst.CSTTransformer):
    def leave_List(self, original_node: cst.List, updated_node: cst.List) -> cst.BaseExpression:
        new_elements = list(updated_node.elements)

        # Determine if the list is multi-line by checking newlines in the elements
        is_multiline = any(
            isinstance(el.comma, cst.Comma) and isinstance(el.comma.whitespace_after, cst.ParenthesizedWhitespace)
            for el in new_elements
        )

        rbracket_comment = None

        # Handle rbracket comment for multi-line lists
        if is_multiline and updated_node.rbracket.whitespace_before:
            rbracket_comment = (
                updated_node.rbracket.whitespace_before.first_line.comment.value
                if updated_node.rbracket.whitespace_before.first_line.comment
                else None
            )

        if rbracket_comment:
            # Move the existing comment to the last element's comma
            last_element = new_elements[-1]
            last_element_comma = last_element.comma or cst.Comma()
            last_element = last_element.with_changes(
                comma=last_element_comma.with_changes(
                    whitespace_after=cst.ParenthesizedWhitespace(
                        first_line=cst.TrailingWhitespace(
                            whitespace=cst.SimpleWhitespace(" "),
                            comment=cst.Comment(rbracket_comment),
                            newline=cst.Newline(),
                        ),
                        indent=True,
                        last_line=cst.SimpleWhitespace("  "),
                    )
                )
            )
            new_elements[-1] = last_element

        # Add the new element
        new_element = cst.Element(
            value=cst.Integer("4"),
            comma=cst.Comma(
                whitespace_after=cst.ParenthesizedWhitespace(
                    first_line=cst.SimpleWhitespace(" "),
                    last_line=cst.SimpleWhitespace(" "),
                )
            ),
        )
        new_elements.append(new_element)

        # Update the rbracket
        updated_rbracket = updated_node.rbracket.with_changes(
            whitespace_before=cst.ParenthesizedWhitespace(
                first_line=cst.TrailingWhitespace(
                    whitespace=cst.SimpleWhitespace(""),
                    comment=cst.Comment("# comment 4"),
                    newline=cst.Newline(),
                ),
                indent=True,
                last_line=cst.SimpleWhitespace(""),
            )
        )

        return updated_node.with_changes(elements=new_elements, rbracket=updated_rbracket)


# Test cases
code_multiline = """[
  1, # comment 1
  2, # comment 2
  3, # comment 3
]"""

code_compact = "[1, 2, 3]"

# Apply the transformation
tree_multiline = cst.parse_module(code_multiline)
modified_tree_multiline = tree_multiline.visit(ListAppender())

tree_compact = cst.parse_module(code_compact)
modified_tree_compact = tree_compact.visit(ListAppender())

print(modified_tree_multiline.code)
print(modified_tree_compact.code)

[dheerajck]

Hi there, I am trying to update an existing code that looks like the example above, is there a better way to do this rn ??
There seems to be many edge cases that I will need to handle if I want to do add a new element with a comment to a list rn.

[rafalkrupinski]

Could you edit the title to make it more descriptive?

I'm also interested in this.

I'd like to generate function definitions and type annotations, but the problem is the same: all code examples are parsed in such way that makes indentations belong to the parent/previous element, instead of the element on the same line. So you need separate indentation for the first element, all elements in the middle and the last element which owns the indentation for the closing parenthesis/bracket