CODEWITHSUNDEEP
pythonabstract-syntax-tree
Asdfg
Asdfg

Reputation: 12213

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"
}

Upvotes: 0

Views: 44

Answers (2)

Frank Yellin
Frank Yellin

Reputation: 11297

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.

Upvotes: 2

Mohit Reddy
Mohit Reddy

Reputation: 141

Well, AST itself is generating the correct structure, its just not formatting it correctly as it doesnt exactly have any options to do that. Instead, you can do some quick post script modification to the dict to get the output you want.

unparsed_code = ast.unparse(my_awesome_module_node)

formatted_code = unparsed_code.replace(
    "{",
    "{\n    "
).replace(
    "}", 
    "\n}"
).replace(
    ", ", 
    ",\n    "
)

print(formatted_code)

This should work for your purposes.

Upvotes: 1

Related Questions