Script to wrap multi-line text blocks in source code dictionary

wrap.py

"""
Script to wrap multi-line text blocks in dictionaries.
Made this for InVEST model ARGS_SPECs "about" properties since I'm 
moving things around and it's tedious to wrap them all by hand.
This could be generalized to any multi-line text block or an entire file.
"""
import ast
import importlib
import textwrap

max_width = 80
tab = '    '

def wrap_str(key, val, indent=0):
    """Wrap a key-value pair from a dictionary.

    Args:
        key (str): the key of the key-value pair. e.g. "about" 
        val (str):
        indent (int): how many tabs to indent the first line by.
            wrapped lines are indented by one additional tab.

    Returns:
        list[str]: list of complete wrapped lines (without newline chars)
    """
    # see if it can all fit in one line
    one_line = f'{tab * indent}"{key}": "{val}"'
    if len(one_line) <= max_width:
        return [one_line]

    wrapper = textwrap.TextWrapper()
    # allow room for a quote on the end (for the wrapped text)
    wrapper.width = max_width - 1  
    # start the first line indented, with key: "...
    wrapper.initial_indent = f'{tab * indent}"{key}": ("'
    # start subsequent lines indented one more level 
    # and with a quote for the wrapped text block
    wrapper.subsequent_indent = f'{tab * (indent + 1)}"'
    # leave whitespace on the ends of wrapped lines
    wrapper.drop_whitespace = False

    lines = wrapper.wrap(val)
    # add a quote to the end of each line of wrapped text
    lines = [line + '"' for line in lines]
    # close the wrapped text block with a parenthesis at the end
    lines[-1] += ')'
    return lines


def wrap_dict(dic, indent=1):

    lines = []
    for key, val in dic.items():
        if type(val) is str:
            lines.append(wrap_str(key, val, indent=indent))
        elif type(val) is dict:
            dict_lines = []
            dict_lines.append(f'{tab * indent}"{key}": {{')
            dict_lines += wrap_dict(val, indent=indent + 1)
            dict_lines.append(tab * indent + '}')
            lines.append(dict_lines)
        else:
            lines.append([f'{tab * indent}"{key}": {val}'])
    # add a comma to the end of all but the last dictionary item
    for i in range(len(lines) - 1):
        lines[i][-1] += ','

    lines = [l for line in lines for l in line]
    return lines


def wrap_dictionary(module, dict_name):
    """Modify source code to correctly wrap a dictionary in place.

    Args:
        module (module): imported module containing the dictionary
        dict_name (str): name of the dictionary variable to wrap

    Returns:
        None
    """
    # parse the python source file into an abstract syntax tree
    with open(module.__file__) as f:
        tree = ast.parse(f.read())

    # find where the variable name is assigned to
    for node in ast.iter_child_nodes(tree):
        if isinstance(node, ast.Assign):
            if node.targets[0].id == name:
                start, end = node.lineno, node.end_lineno
                print(node.lineno, node.end_lineno)

    # read in the text of the file
    with open(module.__file__, newline='') as f:
        lines = f.readlines()
    # identify the linebreak (newline or carriage return)
    # avoid changing the linebreak unnecessarily (minimize diff size)
    if lines[0][-2:] == '\r\n':
            linebreak = '\r\n'
    elif lines[0][-1] == '\n':
        linebreak = '\n'
    else:
        raise ValueError(f'Unknown linebreak: {repr(lines[0][-1])}')

    # wrap each line to the max_width
    wrapped = wrap_dict(getattr(module, name))
    wrapped = [line + linebreak for line in wrapped]

    # print out if any are longer than expected
    for i, line in enumerate(wrapped):
        if len(line) > max_width + len(linebreak):
            print(len(line), line)

    modified = lines[:start - 1]
    modified += [f'{name} = {{{linebreak}']  # add the assignment line back in
    modified += wrapped
    modified += [f'}}{linebreak}']  # final closing bracket for the dict
    modified += lines[end:]

    # overwrite the module file with the modified lines
    with open(module.__file__, 'w', newline='') as f:
        for line in modified:
            f.write(line)


if __name__ == '__main__':
    module = importlib.import_module('src.natcap.invest.carbon')
    wrap_dictionary(module, 'ARGS_SPEC')