Skip to content

Commit a3e9895

Browse files
committed
Address review
1 parent 2772683 commit a3e9895

File tree

3 files changed

+110
-7
lines changed

3 files changed

+110
-7
lines changed

Lib/test/test_ast/test_ast.py

Lines changed: 106 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -222,24 +222,126 @@ def test_negative_locations_for_compile(self):
222222

223223
def test_docstring_optimization_single_node(self):
224224
# https://github.com/python/cpython/issues/137308
225+
class_example1 = textwrap.dedent('''
226+
class A:
227+
"""Docstring"""
228+
''')
229+
class_example2 = textwrap.dedent('''
230+
class A:
231+
"""
232+
Docstring"""
233+
''')
234+
def_example1 = textwrap.dedent('''
235+
def some():
236+
"""Docstring"""
237+
''')
238+
def_example2 = textwrap.dedent('''
239+
def some():
240+
"""Docstring
241+
"""
242+
''')
243+
async_def_example1 = textwrap.dedent('''
244+
async def some():
245+
"""Docstring"""
246+
''')
247+
async_def_example2 = textwrap.dedent('''
248+
async def some():
249+
"""
250+
Docstring
251+
"""
252+
''')
225253
for code in [
226-
'class A: """Docstring"""',
227-
'def func(): """Docstring"""',
228-
'async def func(): """Docstring"""',
254+
class_example1,
255+
class_example2,
256+
def_example1,
257+
def_example2,
258+
async_def_example1,
259+
async_def_example2,
229260
]:
230261
for opt_level in [0, 1, 2]:
231262
with self.subTest(code=code, opt_level=opt_level):
232263
mod = ast.parse(code, optimize=opt_level)
233264
self.assertEqual(len(mod.body[0].body), 1)
234265
if opt_level == 2:
235-
self.assertIsInstance(mod.body[0].body[0], ast.Pass)
266+
pass_stmt = mod.body[0].body[0]
267+
self.assertIsInstance(pass_stmt, ast.Pass)
268+
self.assertEqual(
269+
vars(pass_stmt),
270+
{
271+
'lineno': 3,
272+
'col_offset': 4,
273+
'end_lineno': 3,
274+
'end_col_offset': 8,
275+
},
276+
)
277+
else:
278+
self.assertIsInstance(mod.body[0].body[0], ast.Expr)
279+
self.assertIsInstance(
280+
mod.body[0].body[0].value,
281+
ast.Constant,
282+
)
283+
284+
compile(code, "a", "exec")
285+
compile(code, "a", "exec", optimize=opt_level)
286+
compile(mod, "a", "exec")
287+
compile(mod, "a", "exec", optimize=opt_level)
288+
289+
def test_docstring_optimization_multiple_nodes(self):
290+
# https://github.com/python/cpython/issues/137308
291+
class_example = textwrap.dedent(
292+
"""
293+
class A:
294+
'''
295+
Docstring
296+
'''
297+
x = 1
298+
"""
299+
)
300+
301+
def_example = textwrap.dedent(
302+
"""
303+
def some():
304+
'''
305+
Docstring
306+
307+
'''
308+
x = 1
309+
"""
310+
)
311+
312+
async_def_example = textwrap.dedent(
313+
"""
314+
async def some():
315+
316+
'''Docstring
317+
318+
'''
319+
x = 1
320+
"""
321+
)
322+
323+
for code in [
324+
class_example,
325+
def_example,
326+
async_def_example,
327+
]:
328+
for opt_level in [0, 1, 2]:
329+
with self.subTest(code=code, opt_level=opt_level):
330+
mod = ast.parse(code, optimize=opt_level)
331+
if opt_level == 2:
332+
self.assertNotIsInstance(
333+
mod.body[0].body[0],
334+
(ast.Pass, ast.Expr),
335+
)
236336
else:
237337
self.assertIsInstance(mod.body[0].body[0], ast.Expr)
238338
self.assertIsInstance(
239339
mod.body[0].body[0].value,
240340
ast.Constant,
241341
)
242342

343+
compile(code, "a", "exec")
344+
compile(code, "a", "exec", optimize=opt_level)
243345
compile(mod, "a", "exec")
244346
compile(mod, "a", "exec", optimize=opt_level)
245347

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +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.
1+
A standalone docstring in a node body is optimized as a :keyword:`pass`
2+
statement to ensure that the node's body is never empty. There was a
3+
:exc:`ValueError` in :func:`compile` otherwise.

Python/ast_preprocess.c

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -438,6 +438,7 @@ stmt_seq_remove_item(asdl_stmt_seq *stmts, Py_ssize_t idx)
438438
static int
439439
remove_docstring(asdl_stmt_seq *stmts, Py_ssize_t idx, PyArena *ctx_)
440440
{
441+
assert(_PyAST_GetDocString(stmts) != NULL);
441442
// In case there's just the docstring in the body, replace it with `pass`
442443
// keyword, so body won't be empty.
443444
if (asdl_seq_LEN(stmts) == 1) {

0 commit comments

Comments
 (0)