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

Question

我们正在为 python 模块使用拿破仑风格的文档字符串。但是需要在名为 Data Owner 和 DAL Owner 的文档字符串中自动填充其他属性，以便给定的函数如下所示:

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

给定函数的这些附加属性及其值在单独的 csv 文件中提供。我当时的想法是有一个脚本(awk，sed？)将

• 提取给定 python 文件中的所有函数名称。可以在 python 中轻松完成
• 对于那些函数名称，检查所有者是否存在于 csv 文件中，如果存在，则创建函数名称和所有者的映射。可行

现在，这是我还没有弄清楚并且不知道最好的前进方向的部分。对于给定的函数名称和所有者，我需要返回到 python 文件并将所有者添加到文档字符串(如果存在)。我在考虑某种 awk 脚本，但不太确定

• 找到与模式匹配的函数
• 对于该模式，查看 doctsring 是否存在，右括号后有三重引号
• 如果 docstring 存在，在结束的三重引号之前为所有者添加额外的两行
• 如果 docstring 不存在，则在函数声明后的行中插入两行 for owners between tripe quotations。

我知道这是很多步骤，但任何人都可以提供对前面 4 个要点的洞察力，以便在给定函数、属性和 python 文件的情况下将附加属性插入文档字符串。像 sed、awk 这样的 linux 实用程序会更有用，还是我应该走 python 路线。是否有其他更容易实现的选项。

Answer 1

在 ast 中分配新文档字符串的过程是:

1. 使用 ast.get_docstring 获取现有文档字符串
2. 使用修改后的内容创建一个新的 ast 节点
3. 如果现有的 dostring 是None，则在父节点正文的开头插入新节点
4. 如果有一个现有的文档字符串，用新节点替换它的节点
5. 使用unparse ^* 来自 Cpython Tools 的工具来生成新的源代码(您可能需要从 github 下载这个 - 确保您获得的版本与您的 python 版本相匹配)

下面是一些示例代码:

$  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'

(我自己为 python2.7 回答了类似的问题，here)

^* 从 Python 3.9 开始，ast 模块提供了一个 unparse可以代替 unparse 工具使用的函数:src = ast.unparse(new_tree)