Replacing a SimpleString inside a libcst function definition? (dataclasses.FrozenInstanceError: cannot assign to field 'body')

Question:

Context

While trying to use the libcst module, I am experiencing some difficulties updating a documentation of a function.

MWE

To reproduce the error, the following minimal working example (MWE) is included:

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

Error:

Running python mwe.py yields:

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'

Question

How can one replace the docstring of a function named: add_three in some Python code file_content using the libcst module?

Partial Solution

I found the following solution for a basic example, however, I did not test it on different functions inside classes, with typed arguments, typed returns etc.

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("nnnnModified:")
print(modified_tree.code)

For example, this partial solution fails on:

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

because the solution does not verify the name of the function in which the SimpleString occurs.

Asked By: a.t.

||

Answers:

Why were you getting the "FrozenInstanceError" ?

As you saw, the CST produced by libcst is a graph made of immutable nodes (each one representing a part of the Python language). If you want to change a node, you actually need to make a new copy of it.
This is done using node.with_changes() method.

So you could do this in your first code snippet.
However, there are more "elegant" ways to achieve this, partly documented in libcst’s tutorial, as you just started doing in your partial solution.

How can one replace the docstring of a function named: add_three in some Python code file_content using the libcst module?

Use the libcst.CSTTransformer to navigate your way through:

  1. You need to find, in the CST, the node representing your function (libcst.FunctionDef)
  2. You then need to find the node representing the documentation of your function (licst.SimpleString)
  3. Update this documentation node
import libcst

class DocUpdater(libcst.CSTTransformer):
    """Upodate the docstring of the function `add_three`"""
    def __init__(self) -> None:
        super().__init__()
        self._docstring: str | None = None

    def visit_FunctionDef(self, node: libcst.FunctionDef) -> Optional[bool]:
        """Trying to find the node defining function `add_three`,
         and get its docstring"""
        if node.name.value == 'add_three':
            self._docstring = f'"""{node.get_docstring(clean=False)}"""'
            """Unfortunatly, get_docstring doesn't return the exact docstring
            node value: you need to add the docstring's triple quotes"""
            return True
        return False

    def leave_SimpleString(
        self, original_node: libcst.SimpleString, updated_node: libcst.SimpleString
    ) -> libcst.BaseExpression:
        """Trying to find the node defining the docstring
        of your function, and update the docstring"""
        if original_node.value == self._docstring:
            return updated_node.with_changes(value='"""My new docstring"""')

        return updated_node

And finally:

test = r'''
import foo
from a.b import foo_method


class C:
    def add_three(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
'''

cst = libcst.parse_module(test)
updated_cst = cst.visit(DocUpdater())

print(updated_cst.code)

Output:

import foo
from a.b import foo_method


class C:
    def add_three(self, x):
        """My new docstring"""
        return foo_method(x)

def do_another_thing(y: list[str]) -> int:
    """Bike"""
    return 1
Answered By: zaicruvoir1rominet
Categories: questions Tags: , , ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.