Insert docstring attributes in a python file

Question

We are using napolean style docstring for python modules. But there is a need to auto populate additional attributes in the docstring called Data Owner and DAL Owner so that the given function looks like this:

def func(self, arg1=None, arg2=None):
"""
Returns the timeseries for the specified arg1 and arg2.
Args:
    arg1: argument 1
    arg2: argument 2
Returns:
    DataFrame containing timeseries of arg1 for arg2.

DAL Owner: Team IT
Data Owner: Team A
"""

These additional attributes and their values for a given function are provided in a separate csv file. The way I was thinking was to have a script (awk, sed?) that will

• extract all the function names in a given python file. Can easily do it in python
• for those function names, check if the owners exist in the csv file and if so create a mapping of the function name and owners. Doable

Now, this is the part which I havent figured out and dont know the best way forward. For a given function name and owners, I need to go back into the python file and add the owners to the docstring if it exists. I am thinking some sort of awk script but not quite sure

• Find the function that matches the pattern
• For that pattern, see if doctsring exists, triple quotation marks after closing parenthesis
• If docstring exists, add additional two lines for the owners before the closing triple quotation
• If docstring does not exists, then insert the two lines for owners between tripe quotations on the line after function declaration.

I know this is a lot of steps but can anyone provide insight with the previous 4 bullet points to insert the additional attributes to docstring given the function, attributes and the python file. Will a linux utility like sed, awk be more useful or should I go the python route. Is there some other option that's easier to implement.

Answer 1

The process for assigning a new docstring in an ast is:

1. Get the existing docstring using ast.get_docstring
2. Create a new ast node with amended content
3. If the existing dostring is None, insert the new node at the start of the parent node's body
4. If there was an existing docstring, replace it's node with the new node
5. Use the unparse^* tool from Cpython Tools to generate the new source (you may need to download this from github - ensure you get the version that matches your python version)

Here's some example code:

$  cat fixdocstrings.py                            
import ast                                                                               
import io
from unparse import Unparser


class DocstringWriter(ast.NodeTransformer):
    def visit_FunctionDef(self, node):
        docstring = ast.get_docstring(node)
        new_docstring_node = make_docstring_node(docstring)
        if docstring:
            # Assumes the existing docstring is the first node 
            # in the function body.
            node.body[0] = new_docstring_node
        else:
            node.body.insert(0, new_docstring_node)
        return node


def make_docstring_node(docstring):
    if docstring is None:
        content = "A new docstring"
    else:
        content = docstring + " -- amended"
    s = ast.Str(content)
    return ast.Expr(value=s)


if __name__ == "__main__":
    tree = ast.parse(open("docstringtest.py").read())
    transformer = DocstringWriter()
    new_tree = transformer.visit(tree)
    ast.fix_missing_locations(new_tree)
    buf = io.StringIO()
    Unparser(new_tree, buf)
    buf.seek(0)
    print(buf.read())

$  cat docstringtest.py 
def foo():
    pass


def bar():
    """A docstring."""

$  python fixdocstrings.py 


def foo():
    'A new docstring'
    pass

def bar():
    'A docstring. -- amended'

(I answered something similar for myself for python2.7, here)

^* As of Python 3.9, the ast module provides an unparse function that can be used instead of the unparse tool: src = ast.unparse(new_tree)