Skip to content

Commit 6a1908b

Browse files
Anuoluwapo25SamWilsn
authored andcommitted
chore(tools): clear docstring when creating fork
Closes: #1423
1 parent b6553ba commit 6a1908b

File tree

5 files changed

+149
-3
lines changed

5 files changed

+149
-3
lines changed

src/ethereum_spec_tools/new_fork/builder.py

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -215,7 +215,7 @@ def _to_args(
215215
[
216216
"codemod",
217217
"--no-format",
218-
"string.StringReplaceCommand",
218+
"string_replace.StringReplaceCommand",
219219
]
220220
+ common,
221221
[
@@ -229,6 +229,32 @@ def _to_args(
229229
return commands
230230

231231

232+
@dataclass
233+
class ClearDocstring(CodemodArgs):
234+
"""
235+
Describe how to clear the docstring in __init__.py to libcst.tool:main.
236+
"""
237+
238+
@override
239+
def _to_args(
240+
self, fork_builder: "ForkBuilder", working_directory: Path
241+
) -> list[list[str]]:
242+
init_path = (
243+
working_directory
244+
/ "ethereum"
245+
/ fork_builder.new_fork
246+
/ "__init__.py"
247+
)
248+
return [
249+
[
250+
"codemod",
251+
"remove_docstring.RemoveDocstringCommand",
252+
"--no-format",
253+
str(init_path),
254+
]
255+
]
256+
257+
232258
class ForkBuilder:
233259
"""
234260
Takes a template fork and uses it to generate a new fork, applying source
@@ -353,6 +379,7 @@ def __init__(
353379
RenameFork(),
354380
SetForkCriteria(),
355381
ReplaceForkName(),
382+
ClearDocstring(),
356383
]
357384

358385
def _create_working_directory(self) -> TemporaryDirectory:
Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
"""
2+
libcst codemod that removes the module docstring.
3+
"""
4+
5+
import libcst as cst
6+
from libcst.codemod import CodemodCommand
7+
from typing_extensions import override
8+
9+
10+
class RemoveDocstringCommand(CodemodCommand):
11+
"""
12+
Removes the module docstring if it exists.
13+
"""
14+
15+
DESCRIPTION: str = "Remove the module docstring."
16+
17+
@override
18+
def transform_module_impl(self, tree: cst.Module) -> cst.Module:
19+
"""
20+
Transform the tree by removing the docstring.
21+
"""
22+
if len(tree.body) == 0:
23+
return tree
24+
first_stmt = tree.body[0]
25+
if not isinstance(first_stmt, cst.SimpleStatementLine):
26+
return tree
27+
if len(first_stmt.body) != 1:
28+
return tree
29+
expr = first_stmt.body[0]
30+
if not isinstance(expr, cst.Expr):
31+
return tree
32+
if not isinstance(expr.value, cst.SimpleString):
33+
return tree
34+
new_body = tree.body[1:]
35+
return tree.with_changes(body=new_body)

tests/json_infra/test_tools_new_fork.py

Lines changed: 83 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,15 @@
55
from pathlib import Path
66
from tempfile import TemporaryDirectory
77

8+
import libcst as cst
89
import pytest
10+
from libcst.codemod import CodemodContext
911

1012
from ethereum_spec_tools.forks import Hardfork
1113
from ethereum_spec_tools.new_fork.cli import main as new_fork
14+
from ethereum_spec_tools.new_fork.codemod.remove_docstring import (
15+
RemoveDocstringCommand,
16+
)
1217

1318

1419
@pytest.mark.parametrize(
@@ -56,10 +61,14 @@ def test_end_to_end(template_fork: str) -> None:
5661
with (fork_dir / "__init__.py").open("r") as f:
5762
source = f.read()
5863

64+
assert '"""' not in source[:20]
5965
assert "FORK_CRITERIA = ByTimestamp(7)" in source
60-
assert "E2E Fork" in source
6166
assert template_fork.capitalize() not in source
6267

68+
with (fork_dir / "utils" / "hexadecimal.py").open("r") as f:
69+
source = f.read()
70+
assert "E2E Fork" in source
71+
6372
with (fork_dir / "vm" / "gas.py").open("r") as f:
6473
source = f.read()
6574

@@ -82,3 +91,76 @@ def test_end_to_end(template_fork: str) -> None:
8291
"from ethereum.forks.paris import trie as previous_trie"
8392
in f.read()
8493
)
94+
95+
96+
def has_module_docstring(file_path: Path) -> bool:
97+
"""Return True if the file starts with a module-level doc-string."""
98+
tree = cst.parse_module(file_path.read_text())
99+
if not tree.body:
100+
return False
101+
first = tree.body[0]
102+
if not isinstance(first, cst.SimpleStatementLine):
103+
return False
104+
if len(first.body) != 1:
105+
return False
106+
expr = first.body[0]
107+
return isinstance(expr, cst.Expr) and isinstance(
108+
expr.value, cst.SimpleString
109+
)
110+
111+
112+
def test_remove_docstring_command() -> None:
113+
"""Test that RemoveDocstringCommand removes module docstrings."""
114+
source = '"""Module docstring."""\n\nsome_var = 123\n'
115+
module = cst.parse_module(source)
116+
context = CodemodContext()
117+
command = RemoveDocstringCommand(context)
118+
119+
new_module = command.transform_module(module)
120+
result = new_module.code
121+
122+
assert '"""Module docstring."""' not in result
123+
assert "some_var = 123" in result
124+
125+
126+
def test_remove_docstring_preserves_other_docstrings() -> None:
127+
"""Test that function/class docstrings are preserved."""
128+
source = '''"""Module docstring."""
129+
130+
def foo():
131+
"""Function docstring."""
132+
pass
133+
'''
134+
module = cst.parse_module(source)
135+
context = CodemodContext()
136+
command = RemoveDocstringCommand(context)
137+
138+
new_module = command.transform_module(module)
139+
result = new_module.code
140+
141+
assert not result.startswith('"""Module docstring."""')
142+
assert '"""Function docstring."""' in result
143+
144+
145+
def test_remove_docstring_handles_files_without_docstrings() -> None:
146+
"""Test that files without docstrings remain unchanged."""
147+
source_without_docstring = "some_var = 123\n\ndef foo():\n pass\n"
148+
module = cst.parse_module(source_without_docstring)
149+
context = CodemodContext()
150+
command = RemoveDocstringCommand(context)
151+
152+
new_module = command.transform_module(module)
153+
154+
assert new_module.code == source_without_docstring
155+
156+
157+
def test_remove_docstring_handles_empty_files() -> None:
158+
"""Test that empty files remain empty."""
159+
source_empty = ""
160+
module = cst.parse_module(source_empty)
161+
context = CodemodContext()
162+
command = RemoveDocstringCommand(context)
163+
164+
new_module = command.transform_module(module)
165+
166+
assert new_module.code == source_empty

vulture_whitelist.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,9 @@
3131
from ethereum_spec_tools.lint.lints.import_hygiene import ImportHygiene
3232
from ethereum_spec_tools.new_fork.codemod.comment import CommentReplaceCommand
3333
from ethereum_spec_tools.new_fork.codemod.constant import SetConstantCommand
34-
from ethereum_spec_tools.new_fork.codemod.string import StringReplaceCommand
34+
from ethereum_spec_tools.new_fork.codemod.string_replace import (
35+
StringReplaceCommand,
36+
)
3537

3638
# src/ethereum/utils/hexadecimal.py
3739
hex_to_bytes256

0 commit comments

Comments
 (0)