LibCST example for Moscow Python Conf++ 2021

main.py

RESULTS = {}


def add_and_store(a, b, name):
    """Process and save to DB"""
    result = a + b
    RESULTS[name] = result


def add(a, b):
    return a + b


def store(result, name):
    RESULTS[name] = result


# sum 5 and 10 then store it to database with key 'fifteen'
add_and_store(5, 10, "fifteen")
# another example
add_and_store(2, 12, "value")

import libcst as cst
import libcst.matchers as m

example = cst.parse_expression('store(add(1, 2), "name")')


def print_node(node: cst.CSTNode):
    """Print any CST Node"""
    mod = cst.parse_module("")
    wrapped = mod.with_changes(body=[node])
    print(wrapped.code)


class CodeTransformer(cst.CSTTransformer):
    def visit_Call(self, node: cst.Call):
        """Read-only visitor"""
        if node.func.value == 'add_and_store':
            print_node(node)

    def leave_Expr(self, original_node: cst.Expr, updated_node: cst.Expr) -> cst.CSTNode:
        """Modifying visitor. Must return replacement node"""
        if m.matches(updated_node.value, m.Call(func=m.Name('add_and_store'))):
            args = updated_node.value.args
            add_args = [
                args[0],
                args[1].with_changes(comma=cst.MaybeSentinel.DEFAULT)
            ]
            updated_call = example.with_deep_changes(example.args[0].value, args=add_args)
            store_args = (updated_call.args[0], args[-1])
            updated_call = updated_call.with_changes(args=store_args)
            updated_node = updated_node.with_changes(value=updated_call)
        return updated_node


def main():
    with open('main.py') as codefile:
        code = codefile.read()
    tree = cst.parse_module(code)
    transformer = CodeTransformer()
    modified_code = tree.visit(transformer)
    print(modified_code.code)


if __name__ == '__main__':
    main()