In-place addition of docstrings to a file given to CLI

[robodair]

[pomdtr]

I've done a WIP implementation using parsor.

import parso
from parso.python.tree import DocstringMixin, PythonNode, String
from pydocstring import generate_docstring


def search_successors(root, *node_types):
    successors = []
    def search_successors_recursive(node):
        if node.type in node_types:
            successors.append(node)
        if hasattr(node, 'children'):
            for children in node.children:
                search_successors_recursive(children)

    search_successors_recursive(root)
    return successors


def add_docstring(bloc):
    docstring = '"""' + generate_docstring(bloc.get_code()) + '"""'
    if bloc.type == 'file_input':
        bloc.children.insert(0, PythonNode('simple_stmt', [String(docstring + "\n", (0,0))]))

    if bloc.type in ('classdef', 'funcdef'):
        padding = bloc.start_pos[1] + 4
        docstring = "\n".join([" " * padding + line for line in docstring.splitlines()])

        bloc = bloc.children[bloc.children.index(':') + 1]
        if bloc.type == 'suite':
            bloc.children.insert(1, PythonNode('simple_stmt',[String(docstring + "\n", (0,0))]))

def patch_file(file_content):
    scopes = ('file_input', 'classdef', 'funcdef')
    module = parso.parse(file_content)
    blocs = search_successors(module, *scopes)

    missing_docstring_blocs = [bloc for bloc in blocs if bloc.get_doc_node() is None]
    for bloc in missing_docstring_blocs:
        add_docstring(bloc)
    return module.get_code()


if __name__ == "__main__":
    with open("test.py") as stream:
        file_content = stream.read()

    with open("result.py", "w") as f:
        f.write(patch_file(file_content))

Does it look right to you ?

[robodair]

@pomdtr that looks good to me (sorry I missed a notification on this and didn't reply much earlier). If you would like to make a PR with this and a test that would be awesome!