使用libcst替换指定函数文档字符串遇FrozenInstanceError的解决方法

背景：使用libcst模块时，尝试更新Python函数的文档字符串遇到困难。

最小可复现示例（MWE） #

为复现错误，提供以下最小可复现示例：

from libcst import (  # type: ignore[import]
    Expr,
    FunctionDef,
    IndentedBlock,
    MaybeSentinel,
    SimpleStatementLine,
    SimpleString,
    parse_module,
)

original_content: str = """
\"\"\"Example python file with a function.\"\"\"


from typeguard import typechecked


@typechecked
def add_three(*, x: int) -> int:
    \"\"\"ORIGINAL This is a new docstring core.
    that consists of multiple lines. It also has an empty line inbetween.

    Here is the emtpy line.\"\"\"
    return x + 2

"""
new_docstring_core: str = """\"\"\"This is a new docstring core.
    that consists of multiple lines. It also has an empty line inbetween.

    Here is the emtpy line.\"\"\""""


def replace_docstring(
    original_content: str, func_name: str, new_docstring: str
) -> str:
    """Replaces the docstring in a Python function."""
    module = parse_module(original_content)
    for node in module.body:
        if isinstance(node, FunctionDef) and node.name.value == func_name:
            print("Got function node.")
            # print(f'node.body={node.body}')
            if isinstance(node.body, IndentedBlock):
                if isinstance(node.body.body[0], SimpleStatementLine):
                    simplestatementline: SimpleStatementLine = node.body.body[
                        0
                    ]

                    print("Got SimpleStatementLine")
                    print(f"simplestatementline={simplestatementline}")

                    if isinstance(simplestatementline.body[0], Expr):
                        print(
                            f"simplestatementline.body={simplestatementline.body}"
                        )

                        simplestatementline.body = (
                            Expr(
                                value=SimpleString(
                                    value=new_docstring,
                                    lpar=[],
                                    rpar=[],
                                ),
                                semicolon=MaybeSentinel.DEFAULT,
                            ),
                        )


replace_docstring(
    original_content=original_content,
    func_name="add_three",
    new_docstring=new_docstring_core,
)
print("done")

错误信息 #

运行python mwe.py得到如下错误：

Traceback (most recent call last):
  File "/home/name/git/Hiveminds/jsonmodipy/mwe0.py", line 68, in <module>
    replace_docstring(
  File "/home/name/git/Hiveminds/jsonmodipy/mwe0.py", line 56, in replace_docstring
    simplestatementline.body = (
    ^^^^^^^^^^^^^^^^^^^^^^^^
  File "<string>", line 4, in __setattr__
dataclasses.FrozenInstanceError: cannot assign to field 'body'

问题 #

如何使用libcst模块替换Python代码file_content中名为add_three的函数的文档字符串？

部分解决方案的局限性 #

找到一个基础示例的解决方案，但该方案未考虑类内函数、带类型注解的函数等场景，会无差别替换所有函数的文档字符串。

基础方案代码：

from pprint import pprint
import libcst as cst
import libcst.matchers as m


src = """
import foo
from a.b import foo_method


class C:
    def do_something(self, x):
        \"\"\"Some first line documentation
        Some second line documentation

        Args:something.
        \"\"\"
        return foo_method(x)
"""
new_docstring:str = """\"\"\"THIS IS A NEW DOCSTRING
        Some first line documentation
        Some second line documentation

        Args:somethingSTILLCHANGED.
        \"\"\""""

class ImportFixer(cst.CSTTransformer):
    def leave_SimpleStatementLine(self, orignal_node, updated_node):
        """Replace imports that match our criteria."""
        
        if m.matches(updated_node.body[0], m.Expr()):
            expr=updated_node.body[0]
            if m.matches(expr.value, m.SimpleString()):
                simplestring=expr.value
                print(f'GOTT={simplestring}')
                return updated_node.with_changes(body=[
                    cst.Expr(value=cst.SimpleString(value=new_docstring))
                ])
        return updated_node

source_tree = cst.parse_module(src)
transformer = ImportFixer()
modified_tree = source_tree.visit(transformer)

print("Original:")
print(src)
print("\n\n\n\nModified:")
print(modified_tree.code)

该方案在多函数场景下失效，例如：

src = """
import foo
from a.b import foo_method


class C:
    def do_something(self, x):
        \"\"\"Some first line documentation
        Some second line documentation

        Args:something.
        \"\"\"
        return foo_method(x)
    
def do_another_thing(y:List[str]) -> int:
    \"\"\"Bike\"\"\"
    return 1
"""

原因是该方案未验证SimpleString所属的函数名称，会同时替换do_something和do_another_thing的文档字符串。

正确解决方案 #

LibCST的节点是不可变对象，不能直接修改属性，必须使用with_changes方法创建新节点。同时，我们需要跟踪当前遍历到的函数名称，确保只替换目标函数的文档字符串。以下是支持模块级函数、类内函数的通用解决方案：

import libcst as cst
import libcst.matchers as m

def replace_function_docstring(source_code: str, target_func: str, new_docstring: str) -> str:
    class DocstringReplacer(cst.CSTTransformer):
        def __init__(self, target_func: str):
            self.target_func = target_func
            self.current_func = None

        # 进入函数定义时记录当前函数名称
        def visit_FunctionDef(self, node: cst.FunctionDef) -> None:
            self.current_func = node.name.value

        # 离开函数定义时重置当前函数名称
        def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef:
            self.current_func = None
            return updated_node

        def leave_SimpleStatementLine(self, original_node: cst.SimpleStatementLine, updated_node: cst.SimpleStatementLine) -> cst.SimpleStatementLine:
            # 仅当当前处于目标函数内，且节点是文档字符串时替换
            if self.current_func == self.target_func:
                if m.matches(updated_node.body[0], m.Expr(value=m.SimpleString())):
                    # 创建新的文档字符串节点
                    new_expr = cst.Expr(value=cst.SimpleString(value=new_docstring))
                    return updated_node.with_changes(body=[new_expr])
            return updated_node

    # 解析源码并应用转换器
    module = cst.parse_module(source_code)
    replacer = DocstringReplacer(target_func)
    modified_module = module.visit(replacer)
    return modified_module.code

# 测试示例
if __name__ == "__main__":
    original_content = """
\"\"\"Example python file with a function.\"\"\"


from typeguard import typechecked


@typechecked
def add_three(*, x: int) -> int:
    \"\"\"ORIGINAL This is a new docstring core.
    that consists of multiple lines. It also has an empty line inbetween.

    Here is the emtpy line.\"\"\"
    return x + 2

class Calculator:
    def multiply(self, a: int, b: int) -> int:
        \"\"\"Multiply two integers.\"\"\"
        return a * b
"""
    new_doc = """\"\"\"Adds 3 to the given integer.

Args:
    x: The integer to add 3 to.

Returns:
    The result of x + 3.\"\"\""""

    modified_code = replace_function_docstring(original_content, "add_three", new_doc)
    print(modified_code)

方案说明 #

1. 不可变节点处理：所有节点修改都通过with_changes方法创建新节点，避免直接赋值引发的FrozenInstanceError。
2. 跟踪当前函数：通过visit_FunctionDef和leave_FunctionDef记录当前遍历的函数名称，确保只替换目标函数的文档字符串。
3. 通用场景支持：同时支持模块级函数和类内的成员函数，兼容带类型注解、装饰器的函数。

内容的提问来源于stack exchange，提问作者a.t.