55from pathlib import Path
66from tempfile import TemporaryDirectory
77
8+ import libcst as cst
89import pytest
10+ from libcst .codemod import CodemodContext
911
1012from ethereum_spec_tools .forks import Hardfork
1113from 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 \n some_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 \n def 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
0 commit comments