Search code examples
pythonabstract-syntax-tree

how to format output with newlines using ast module


I am trying to generate some code using AST module. This is how my code looks like:

I would like to add new line after each key with proper formatting so that the output looks formatted.

my_awesome_properties = [
    (ast.Constant(value="key1"), ast.Constant(value="value1")),
    (ast.Constant(value="key2"), ast.Constant(value="value2")),
    (ast.Constant(value="key3"), ast.Constant(value="value3")),
]

my_awesome_properties_dict = ast.Dict(
    keys=[key for key, value in my_awesome_properties],
    values=[value for key, value in my_awesome_properties]
)

my_awesome_properties_assignment = ast.Assign(
    targets=[ast.Name(id='my_awesome_properties', ctx=ast.Store())],
    value=my_awesome_properties_dict
)

my_awesome_properties_assignment.lineno = 1


my_awesome_module_node = ast.Module(
    body=[
        my_awesome_properties_assignment,
    ],
    type_ignores=[]
)

print(ast.unparse(my_awesome_module_node))

This generates output as my_awesome_properties = {'key1': 'value1', 'key2': 'value2', 'key3': 'value3'}

How do I add new line after each key with proper formatting so that the output looks like this with double quotes?

my_awesome_properties = {
    "key1": "value1",
    "key2": "value2",
    "key3": "value3"
}

Solution

  • You might consider taking your string output and passing it to black.

    import black
    
    formatted_code = black.format_str(code, mode=black.FileMode())
    

    I don't think that AST supports what you want directly.