While trying to use the libcst module, I am experiencing some difficulties updating a documentation of a function.
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")
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'
How can one replace the docstring of a function named: add_three in some Python code file_content using the libcst module?
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("\n\n\n\nModified:")
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.
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:
libcst.FunctionDef)licst.SimpleString)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