# This also must not crash:
ast.parse(tree, optimize=2)
+ def test_docstring_optimization_single_node(self):
+ # https://github.com/python/cpython/issues/137308
+ class_example1 = textwrap.dedent('''
+ class A:
+ """Docstring"""
+ ''')
+ class_example2 = textwrap.dedent('''
+ class A:
+ """
+ Docstring"""
+ ''')
+ def_example1 = textwrap.dedent('''
+ def some():
+ """Docstring"""
+ ''')
+ def_example2 = textwrap.dedent('''
+ def some():
+ """Docstring
+ """
+ ''')
+ async_def_example1 = textwrap.dedent('''
+ async def some():
+ """Docstring"""
+ ''')
+ async_def_example2 = textwrap.dedent('''
+ async def some():
+ """
+ Docstring
+ """
+ ''')
+ for code in [
+ class_example1,
+ class_example2,
+ def_example1,
+ def_example2,
+ async_def_example1,
+ async_def_example2,
+ ]:
+ for opt_level in [0, 1, 2]:
+ with self.subTest(code=code, opt_level=opt_level):
+ mod = ast.parse(code, optimize=opt_level)
+ self.assertEqual(len(mod.body[0].body), 1)
+ if opt_level == 2:
+ pass_stmt = mod.body[0].body[0]
+ self.assertIsInstance(pass_stmt, ast.Pass)
+ self.assertEqual(
+ vars(pass_stmt),
+ {
+ 'lineno': 3,
+ 'col_offset': 4,
+ 'end_lineno': 3,
+ 'end_col_offset': 8,
+ },
+ )
+ else:
+ self.assertIsInstance(mod.body[0].body[0], ast.Expr)
+ self.assertIsInstance(
+ mod.body[0].body[0].value,
+ ast.Constant,
+ )
+
+ compile(code, "a", "exec")
+ compile(code, "a", "exec", optimize=opt_level)
+ compile(mod, "a", "exec")
+ compile(mod, "a", "exec", optimize=opt_level)
+
+ def test_docstring_optimization_multiple_nodes(self):
+ # https://github.com/python/cpython/issues/137308
+ class_example = textwrap.dedent(
+ """
+ class A:
+ '''
+ Docstring
+ '''
+ x = 1
+ """
+ )
+
+ def_example = textwrap.dedent(
+ """
+ def some():
+ '''
+ Docstring
+
+ '''
+ x = 1
+ """
+ )
+
+ async_def_example = textwrap.dedent(
+ """
+ async def some():
+
+ '''Docstring
+
+ '''
+ x = 1
+ """
+ )
+
+ for code in [
+ class_example,
+ def_example,
+ async_def_example,
+ ]:
+ for opt_level in [0, 1, 2]:
+ with self.subTest(code=code, opt_level=opt_level):
+ mod = ast.parse(code, optimize=opt_level)
+ if opt_level == 2:
+ self.assertNotIsInstance(
+ mod.body[0].body[0],
+ (ast.Pass, ast.Expr),
+ )
+ else:
+ self.assertIsInstance(mod.body[0].body[0], ast.Expr)
+ self.assertIsInstance(
+ mod.body[0].body[0].value,
+ ast.Constant,
+ )
+
+ compile(code, "a", "exec")
+ compile(code, "a", "exec", optimize=opt_level)
+ compile(mod, "a", "exec")
+ compile(mod, "a", "exec", optimize=opt_level)
+
def test_slice(self):
slc = ast.parse("x[::]").body[0].value.slice
self.assertIsNone(slc.upper)
return 1;
}
+static int
+remove_docstring(asdl_stmt_seq *stmts, Py_ssize_t idx, PyArena *ctx_)
+{
+ assert(_PyAST_GetDocString(stmts) != NULL);
+ // In case there's just the docstring in the body, replace it with `pass`
+ // keyword, so body won't be empty.
+ if (asdl_seq_LEN(stmts) == 1) {
+ stmt_ty docstring = (stmt_ty)asdl_seq_GET(stmts, 0);
+ stmt_ty pass = _PyAST_Pass(
+ docstring->lineno, docstring->col_offset,
+ // we know that `pass` always takes 4 chars and a single line,
+ // while docstring can span on multiple lines
+ docstring->lineno, docstring->col_offset + 4,
+ ctx_
+ );
+ if (pass == NULL) {
+ return 0;
+ }
+ asdl_seq_SET(stmts, 0, pass);
+ return 1;
+ }
+ // In case there are more than 1 body items, just remove the docstring.
+ return stmt_seq_remove_item(stmts, idx);
+}
+
static int
astfold_body(asdl_stmt_seq *stmts, PyArena *ctx_, _PyASTPreprocessState *state)
{
int docstring = _PyAST_GetDocString(stmts) != NULL;
if (docstring && (state->optimize >= 2)) {
/* remove the docstring */
- if (!stmt_seq_remove_item(stmts, 0)) {
+ if (!remove_docstring(stmts, 0, ctx_)) {
return 0;
}
docstring = 0;