Skip to content

Commit 2772683

Browse files
committed
gh-137308: Replace a single docstring with pass
1 parent fe0e921 commit 2772683

File tree

3 files changed

+51
-1
lines changed

3 files changed

+51
-1
lines changed

Lib/test/test_ast/test_ast.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -220,6 +220,29 @@ def test_negative_locations_for_compile(self):
220220
# This also must not crash:
221221
ast.parse(tree, optimize=2)
222222

223+
def test_docstring_optimization_single_node(self):
224+
# https://github.com/python/cpython/issues/137308
225+
for code in [
226+
'class A: """Docstring"""',
227+
'def func(): """Docstring"""',
228+
'async def func(): """Docstring"""',
229+
]:
230+
for opt_level in [0, 1, 2]:
231+
with self.subTest(code=code, opt_level=opt_level):
232+
mod = ast.parse(code, optimize=opt_level)
233+
self.assertEqual(len(mod.body[0].body), 1)
234+
if opt_level == 2:
235+
self.assertIsInstance(mod.body[0].body[0], ast.Pass)
236+
else:
237+
self.assertIsInstance(mod.body[0].body[0], ast.Expr)
238+
self.assertIsInstance(
239+
mod.body[0].body[0].value,
240+
ast.Constant,
241+
)
242+
243+
compile(mod, "a", "exec")
244+
compile(mod, "a", "exec", optimize=opt_level)
245+
223246
def test_slice(self):
224247
slc = ast.parse("x[::]").body[0].value.slice
225248
self.assertIsNone(slc.upper)
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
Replace a single docstring in a node body to a :keyword:`pass`, so node's
2+
body is never empty. There was a :exc:`ValueError` in :func:`compile`
3+
otherwise.

Python/ast_preprocess.c

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -435,13 +435,37 @@ stmt_seq_remove_item(asdl_stmt_seq *stmts, Py_ssize_t idx)
435435
return 1;
436436
}
437437

438+
static int
439+
remove_docstring(asdl_stmt_seq *stmts, Py_ssize_t idx, PyArena *ctx_)
440+
{
441+
// In case there's just the docstring in the body, replace it with `pass`
442+
// keyword, so body won't be empty.
443+
if (asdl_seq_LEN(stmts) == 1) {
444+
stmt_ty docstring = (stmt_ty)asdl_seq_GET(stmts, 0);
445+
stmt_ty pass = _PyAST_Pass(
446+
docstring->lineno, docstring->col_offset,
447+
// we know that `pass` always takes 4 chars and a single line,
448+
// while docstring can span on multiple lines
449+
docstring->lineno, docstring->col_offset + 4,
450+
ctx_
451+
);
452+
if (pass == NULL) {
453+
return 0;
454+
}
455+
asdl_seq_SET(stmts, 0, pass);
456+
return 1;
457+
}
458+
// In case there are more than 1 body items, just remove the docstring.
459+
return stmt_seq_remove_item(stmts, idx);
460+
}
461+
438462
static int
439463
astfold_body(asdl_stmt_seq *stmts, PyArena *ctx_, _PyASTPreprocessState *state)
440464
{
441465
int docstring = _PyAST_GetDocString(stmts) != NULL;
442466
if (docstring && (state->optimize >= 2)) {
443467
/* remove the docstring */
444-
if (!stmt_seq_remove_item(stmts, 0)) {
468+
if (!remove_docstring(stmts, 0, ctx_)) {
445469
return 0;
446470
}
447471
docstring = 0;

0 commit comments

Comments
 (0)