在 python 文件中插入文档字符串属性

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)