Instagram / LibCST

A concrete syntax tree parser and serializer library for Python that preserves many aspects of Python's abstract syntax tree
https://libcst.readthedocs.io/
Other
1.56k stars 192 forks source link

Add statement to class #1201

Open yowoda opened 2 months ago

yowoda commented 2 months ago

I can't figure out how to add a statetement parsed using libcst.parse_statement to the body of a libcst.ClassDef node in leave_ClassDef. Is there any easy way to implement this?

yowoda commented 2 months ago
def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef) -> cst.ClassDef:
    __class_getitem__Node = cst.parse_statement(
        "__class_getitem__ = classmethod(types.GenericAlias)"
    )
    statements = original_node.body.body
    if isinstance(original_node.body, cst.SimpleStatementSuite):
        statements = [cst.SimpleStatementLine(statements)]

    return updated_node.with_changes(
        body=cst.IndentedBlock(
            body=(
                *statements,
                __class_getitem__Node
            )
        )
    )

Ok this seems to be working, however it would be great if there could be helper functions for this so you don't have to manually check whether the class only consists of smth like class A: pass or class A: ..., which automatically appends a statement to the class body