python-3.xabstract-syntax-treelibcst

Python Libcst: Unable to generate code from the node in visitor class


I am searching for a way to get code from the node in the visitor. Example:

import libcst

code_example = """
from ast import parse
threshold  =  1
print(threshold)
"""

class CodeVisitor(libcst.CSTVisitor):
    def visit_Assign(self, node: libcst.Assign) -> bool:
        print(node)
        return True


demo = libcst.parse_module(code_example)
demo.visit(CodeVisitor())

In the above code I want to get the code(i.e. threshold = 1) of the node . But it seems like libcst doesn't provide that support. I further looked around and figured out a function name code_for_node(node: libcst._nodes.base.CSTNode) → str libcst.Module.code_for_node which belongs to the Module . But I didn't able to find enough help to use this in my code.

Looking forward for the help. Thanks in advance!


Solution

  • After spending sometime, I figured out the way to solve the problem. Here is the code.

    Code:

    import libcst as cst
    
    code_example = """
    from ast import parse
    threshold  =  1
    print(threshold)
    """
    
    class CodeVisitor(cst.CSTVisitor):
        def visit_Assign(self, node: cst.Assign) -> bool:
            print("--> NODE TREE: \n{}".format(node))
            print("--> CODE LINE FROM NODE TREE: \n{}".format(cst.parse_module("").code_for_node(node)))
            return True
    
    
    demo = cst.parse_module(code_example)
    _ = demo.visit(CodeVisitor())
    

    Output:

    --> NODE TREE: 
    Assign(
        targets=[
            AssignTarget(
                target=Name(
                    value='threshold',
                    lpar=[],
                    rpar=[],
                ),
                whitespace_before_equal=SimpleWhitespace(
                    value='  ',
                ),
                whitespace_after_equal=SimpleWhitespace(
                    value='  ',
                ),
            ),
        ],
        value=Integer(
            value='1',
            lpar=[],
            rpar=[],
        ),
        semicolon=MaybeSentinel.DEFAULT,
    )
    --> CODE LINE FROM NODE TREE: 
    threshold  =  1