pythongh-129157: Change the location of with statement AST nodes to spawn the first line

Signed-off-by: Pablo Galindo <[email protected]>

Author: pablogsal

Grammar/python.gram

@@ -394,14 +394,14 @@ for_stmt[stmt_ty]:
 
 with_stmt[stmt_ty]:
     | invalid_with_stmt_indent
-    | 'with' '(' a[asdl_withitem_seq*]=','.with_item+ ','? ')' ':' tc=[TYPE_COMMENT] b=block {
-       _PyAST_With(a, b, NEW_TYPE_COMMENT(p, tc), EXTRA) }
-    | 'with' a[asdl_withitem_seq*]=','.with_item+ ':' tc=[TYPE_COMMENT] b=block {
-        _PyAST_With(a, b, NEW_TYPE_COMMENT(p, tc), EXTRA) }
-    | 'async' 'with' '(' a[asdl_withitem_seq*]=','.with_item+ ','? ')' ':' b=block {
-       CHECK_VERSION(stmt_ty, 5, "Async with statements are", _PyAST_AsyncWith(a, b, NULL, EXTRA)) }
-    | 'async' 'with' a[asdl_withitem_seq*]=','.with_item+ ':' tc=[TYPE_COMMENT] b=block {
-       CHECK_VERSION(stmt_ty, 5, "Async with statements are", _PyAST_AsyncWith(a, b, NEW_TYPE_COMMENT(p, tc), EXTRA)) }
+    | start='with' '(' a[asdl_withitem_seq*]=','.with_item+ ','? ')' end=':' tc=[TYPE_COMMENT] b=block {
+       _PyAST_With(a, b, NEW_TYPE_COMMENT(p, tc), HEADER(start, end), EXTRA) }
+    | start='with' a[asdl_withitem_seq*]=','.with_item+ end=':' tc=[TYPE_COMMENT] b=block {
+        _PyAST_With(a, b, NEW_TYPE_COMMENT(p, tc), HEADER(start, end), EXTRA) }
+    | start='async' 'with' '(' a[asdl_withitem_seq*]=','.with_item+ ','? ')' end=':' b=block {
+       CHECK_VERSION(stmt_ty, 5, "Async with statements are", _PyAST_AsyncWith(a, b, NULL, HEADER(start, end), EXTRA)) }
+    | start='async' 'with' a[asdl_withitem_seq*]=','.with_item+ end=':' tc=[TYPE_COMMENT] b=block {
+       CHECK_VERSION(stmt_ty, 5, "Async with statements are", _PyAST_AsyncWith(a, b, NEW_TYPE_COMMENT(p, tc), HEADER(start, end), EXTRA)) }
     | invalid_with_stmt
 
 with_item[withitem_ty]:

Include/internal/pycore_ast.h

@@ -49,12 +49,12 @@ Module(body=[If(test=Name(id='a', ctx=Load(...)), body=[Pass()], orelse=[If(test
 Module(body=[If(test=Name(id='a', ctx=Load(...)), body=[Pass()], orelse=[Pass()])], type_ignores=[])
 Module(body=[If(test=Name(id='a', ctx=Load(...)), body=[Pass()], orelse=[If(test=Name(...), body=[Pass(...)], orelse=[Pass(...)])])], type_ignores=[])
 Module(body=[If(test=Name(id='a', ctx=Load(...)), body=[Pass()], orelse=[If(test=Name(...), body=[Pass(...)], orelse=[If(...)])])], type_ignores=[])
-Module(body=[With(items=[withitem(context_expr=Name(...), optional_vars=None)], body=[Pass()], type_comment=None)], type_ignores=[])
-Module(body=[With(items=[withitem(context_expr=Name(...), optional_vars=None), withitem(context_expr=Name(...), optional_vars=None)], body=[Pass()], type_comment=None)], type_ignores=[])
-Module(body=[With(items=[withitem(context_expr=Name(...), optional_vars=Name(...))], body=[Pass()], type_comment=None)], type_ignores=[])
-Module(body=[With(items=[withitem(context_expr=Name(...), optional_vars=Name(...)), withitem(context_expr=Name(...), optional_vars=Name(...))], body=[Pass()], type_comment=None)], type_ignores=[])
-Module(body=[With(items=[withitem(context_expr=Name(...), optional_vars=Name(...))], body=[Pass()], type_comment=None)], type_ignores=[])
-Module(body=[With(items=[withitem(context_expr=Name(...), optional_vars=None), withitem(context_expr=Name(...), optional_vars=None)], body=[Pass()], type_comment=None)], type_ignores=[])
+Module(body=[With(items=[withitem(context_expr=Name(...), optional_vars=None)], body=[Pass()], type_comment=None, head=header())], type_ignores=[])
+Module(body=[With(items=[withitem(context_expr=Name(...), optional_vars=None), withitem(context_expr=Name(...), optional_vars=None)], body=[Pass()], type_comment=None, head=header())], type_ignores=[])
+Module(body=[With(items=[withitem(context_expr=Name(...), optional_vars=Name(...))], body=[Pass()], type_comment=None, head=header())], type_ignores=[])
+Module(body=[With(items=[withitem(context_expr=Name(...), optional_vars=Name(...)), withitem(context_expr=Name(...), optional_vars=Name(...))], body=[Pass()], type_comment=None, head=header())], type_ignores=[])
+Module(body=[With(items=[withitem(context_expr=Name(...), optional_vars=Name(...))], body=[Pass()], type_comment=None, head=header())], type_ignores=[])
+Module(body=[With(items=[withitem(context_expr=Name(...), optional_vars=None), withitem(context_expr=Name(...), optional_vars=None)], body=[Pass()], type_comment=None, head=header())], type_ignores=[])
 Module(body=[Raise(exc=None, cause=None)], type_ignores=[])
 Module(body=[Raise(exc=Call(func=Name(...), args=[Constant(...)], keywords=[]), cause=None)], type_ignores=[])
 Module(body=[Raise(exc=Name(id='Exception', ctx=Load(...)), cause=None)], type_ignores=[])
@@ -88,7 +88,7 @@ Module(body=[Expr(value=SetComp(elt=Name(...), generators=[comprehension(...)]))
 Module(body=[Expr(value=SetComp(elt=Name(...), generators=[comprehension(...)]))], type_ignores=[])
 Module(body=[AsyncFunctionDef(name='f', args=arguments(posonlyargs=[], args=[], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[]), body=[Expr(value=Constant(...)), Expr(value=Await(...))], decorator_list=[], returns=None, type_comment=None, type_params=[])], type_ignores=[])
 Module(body=[AsyncFunctionDef(name='f', args=arguments(posonlyargs=[], args=[], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[]), body=[AsyncFor(target=Name(...), iter=Name(...), body=[Expr(...)], orelse=[Expr(...)], type_comment=None)], decorator_list=[], returns=None, type_comment=None, type_params=[])], type_ignores=[])
-Module(body=[AsyncFunctionDef(name='f', args=arguments(posonlyargs=[], args=[], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[]), body=[AsyncWith(items=[withitem(...)], body=[Expr(...)], type_comment=None)], decorator_list=[], returns=None, type_comment=None, type_params=[])], type_ignores=[])
+Module(body=[AsyncFunctionDef(name='f', args=arguments(posonlyargs=[], args=[], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[]), body=[AsyncWith(items=[withitem(...)], body=[Expr(...)], type_comment=None, head=header(...))], decorator_list=[], returns=None, type_comment=None, type_params=[])], type_ignores=[])
 Module(body=[Expr(value=Dict(keys=[None, Constant(...)], values=[Dict(...), Constant(...)]))], type_ignores=[])
 Module(body=[Expr(value=Set(elts=[Starred(...), Constant(...)]))], type_ignores=[])
 Module(body=[FunctionDef(name='f', args=arguments(posonlyargs=[], args=[], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[]), body=[Expr(value=Yield(...))], decorator_list=[], returns=None, type_comment=None, type_params=[])], type_ignores=[])

Include/internal/pycore_ast_state.h

@@ -434,12 +434,12 @@ def main():
 ('Module', [('If', (1, 0, 4, 6), ('Name', (1, 3, 1, 4), 'a', ('Load',)), [('Pass', (2, 2, 2, 6))], [('Pass', (4, 2, 4, 6))])], []),
 ('Module', [('If', (1, 0, 6, 6), ('Name', (1, 3, 1, 4), 'a', ('Load',)), [('Pass', (2, 2, 2, 6))], [('If', (3, 0, 6, 6), ('Name', (3, 5, 3, 6), 'b', ('Load',)), [('Pass', (4, 2, 4, 6))], [('Pass', (6, 2, 6, 6))])])], []),
 ('Module', [('If', (1, 0, 10, 6), ('Name', (1, 3, 1, 4), 'a', ('Load',)), [('Pass', (2, 2, 2, 6))], [('If', (3, 0, 10, 6), ('Name', (3, 5, 3, 6), 'b', ('Load',)), [('Pass', (4, 2, 4, 6))], [('If', (5, 0, 10, 6), ('Name', (5, 5, 5, 6), 'b', ('Load',)), [('Pass', (6, 2, 6, 6))], [('If', (7, 0, 10, 6), ('Name', (7, 5, 7, 6), 'b', ('Load',)), [('Pass', (8, 2, 8, 6))], [('Pass', (10, 2, 10, 6))])])])])], []),
-('Module', [('With', (1, 0, 1, 12), [('withitem', ('Name', (1, 5, 1, 6), 'x', ('Load',)), None)], [('Pass', (1, 8, 1, 12))], None)], []),
-('Module', [('With', (1, 0, 1, 15), [('withitem', ('Name', (1, 5, 1, 6), 'x', ('Load',)), None), ('withitem', ('Name', (1, 8, 1, 9), 'y', ('Load',)), None)], [('Pass', (1, 11, 1, 15))], None)], []),
-('Module', [('With', (1, 0, 1, 17), [('withitem', ('Name', (1, 5, 1, 6), 'x', ('Load',)), ('Name', (1, 10, 1, 11), 'y', ('Store',)))], [('Pass', (1, 13, 1, 17))], None)], []),
-('Module', [('With', (1, 0, 1, 25), [('withitem', ('Name', (1, 5, 1, 6), 'x', ('Load',)), ('Name', (1, 10, 1, 11), 'y', ('Store',))), ('withitem', ('Name', (1, 13, 1, 14), 'z', ('Load',)), ('Name', (1, 18, 1, 19), 'q', ('Store',)))], [('Pass', (1, 21, 1, 25))], None)], []),
-('Module', [('With', (1, 0, 1, 19), [('withitem', ('Name', (1, 6, 1, 7), 'x', ('Load',)), ('Name', (1, 11, 1, 12), 'y', ('Store',)))], [('Pass', (1, 15, 1, 19))], None)], []),
-('Module', [('With', (1, 0, 1, 17), [('withitem', ('Name', (1, 6, 1, 7), 'x', ('Load',)), None), ('withitem', ('Name', (1, 9, 1, 10), 'y', ('Load',)), None)], [('Pass', (1, 13, 1, 17))], None)], []),
+('Module', [('With', (1, 0, 1, 12), [('withitem', ('Name', (1, 5, 1, 6), 'x', ('Load',)), None)], [('Pass', (1, 8, 1, 12))], None, ('header', (1, 0, 1, 7)))], []),
+('Module', [('With', (1, 0, 1, 15), [('withitem', ('Name', (1, 5, 1, 6), 'x', ('Load',)), None), ('withitem', ('Name', (1, 8, 1, 9), 'y', ('Load',)), None)], [('Pass', (1, 11, 1, 15))], None, ('header', (1, 0, 1, 10)))], []),
+('Module', [('With', (1, 0, 1, 17), [('withitem', ('Name', (1, 5, 1, 6), 'x', ('Load',)), ('Name', (1, 10, 1, 11), 'y', ('Store',)))], [('Pass', (1, 13, 1, 17))], None, ('header', (1, 0, 1, 12)))], []),
+('Module', [('With', (1, 0, 1, 25), [('withitem', ('Name', (1, 5, 1, 6), 'x', ('Load',)), ('Name', (1, 10, 1, 11), 'y', ('Store',))), ('withitem', ('Name', (1, 13, 1, 14), 'z', ('Load',)), ('Name', (1, 18, 1, 19), 'q', ('Store',)))], [('Pass', (1, 21, 1, 25))], None, ('header', (1, 0, 1, 20)))], []),
+('Module', [('With', (1, 0, 1, 19), [('withitem', ('Name', (1, 6, 1, 7), 'x', ('Load',)), ('Name', (1, 11, 1, 12), 'y', ('Store',)))], [('Pass', (1, 15, 1, 19))], None, ('header', (1, 0, 1, 14)))], []),
+('Module', [('With', (1, 0, 1, 17), [('withitem', ('Name', (1, 6, 1, 7), 'x', ('Load',)), None), ('withitem', ('Name', (1, 9, 1, 10), 'y', ('Load',)), None)], [('Pass', (1, 13, 1, 17))], None, ('header', (1, 0, 1, 12)))], []),
 ('Module', [('Raise', (1, 0, 1, 5), None, None)], []),
 ('Module', [('Raise', (1, 0, 1, 25), ('Call', (1, 6, 1, 25), ('Name', (1, 6, 1, 15), 'Exception', ('Load',)), [('Constant', (1, 16, 1, 24), 'string', None)], []), None)], []),
 ('Module', [('Raise', (1, 0, 1, 15), ('Name', (1, 6, 1, 15), 'Exception', ('Load',)), None)], []),
@@ -473,7 +473,7 @@ def main():
 ('Module', [('Expr', (1, 0, 1, 16), ('SetComp', (1, 0, 1, 16), ('Name', (1, 1, 1, 2), 'r', ('Load',)), [('comprehension', ('Tuple', (1, 7, 1, 10), [('Name', (1, 7, 1, 8), 'l', ('Store',)), ('Name', (1, 9, 1, 10), 'm', ('Store',))], ('Store',)), ('Name', (1, 14, 1, 15), 'x', ('Load',)), [], 0)]))], []),
 ('Module', [('AsyncFunctionDef', (1, 0, 3, 18), 'f', ('arguments', [], [], None, [], [], None, []), [('Expr', (2, 1, 2, 17), ('Constant', (2, 1, 2, 17), 'async function', None)), ('Expr', (3, 1, 3, 18), ('Await', (3, 1, 3, 18), ('Call', (3, 7, 3, 18), ('Name', (3, 7, 3, 16), 'something', ('Load',)), [], [])))], [], None, None, [])], []),
 ('Module', [('AsyncFunctionDef', (1, 0, 3, 8), 'f', ('arguments', [], [], None, [], [], None, []), [('AsyncFor', (2, 1, 3, 8), ('Name', (2, 11, 2, 12), 'e', ('Store',)), ('Name', (2, 16, 2, 17), 'i', ('Load',)), [('Expr', (2, 19, 2, 20), ('Constant', (2, 19, 2, 20), 1, None))], [('Expr', (3, 7, 3, 8), ('Constant', (3, 7, 3, 8), 2, None))], None)], [], None, None, [])], []),
-('Module', [('AsyncFunctionDef', (1, 0, 2, 21), 'f', ('arguments', [], [], None, [], [], None, []), [('AsyncWith', (2, 1, 2, 21), [('withitem', ('Name', (2, 12, 2, 13), 'a', ('Load',)), ('Name', (2, 17, 2, 18), 'b', ('Store',)))], [('Expr', (2, 20, 2, 21), ('Constant', (2, 20, 2, 21), 1, None))], None)], [], None, None, [])], []),
+('Module', [('AsyncFunctionDef', (1, 0, 2, 21), 'f', ('arguments', [], [], None, [], [], None, []), [('AsyncWith', (2, 1, 2, 21), [('withitem', ('Name', (2, 12, 2, 13), 'a', ('Load',)), ('Name', (2, 17, 2, 18), 'b', ('Store',)))], [('Expr', (2, 20, 2, 21), ('Constant', (2, 20, 2, 21), 1, None))], None, ('header', (2, 1, 2, 19)))], [], None, None, [])], []),
 ('Module', [('Expr', (1, 0, 1, 14), ('Dict', (1, 0, 1, 14), [None, ('Constant', (1, 10, 1, 11), 2, None)], [('Dict', (1, 3, 1, 8), [('Constant', (1, 4, 1, 5), 1, None)], [('Constant', (1, 6, 1, 7), 2, None)]), ('Constant', (1, 12, 1, 13), 3, None)]))], []),
 ('Module', [('Expr', (1, 0, 1, 12), ('Set', (1, 0, 1, 12), [('Starred', (1, 1, 1, 8), ('Set', (1, 2, 1, 8), [('Constant', (1, 3, 1, 4), 1, None), ('Constant', (1, 6, 1, 7), 2, None)]), ('Load',)), ('Constant', (1, 10, 1, 11), 3, None)]))], []),
 ('Module', [('FunctionDef', (1, 0, 1, 16), 'f', ('arguments', [], [], None, [], [], None, []), [('Expr', (1, 9, 1, 16), ('Yield', (1, 9, 1, 16), ('Constant', (1, 15, 1, 16), 1, None)))], [], None, None, [])], []),

Lib/test/test_ast/data/ast_repr.txt

@@ -2267,8 +2267,8 @@ def f():
             start_line, end_line, start_col, end_col = instr.positions
             self.assertEqual(start_line, f.__code__.co_firstlineno + 1)
             self.assertEqual(end_line, f.__code__.co_firstlineno + 1)
-            self.assertEqual(start_col, 17)
-            self.assertEqual(end_col, 20)
+            self.assertEqual(start_col, 12)
+            self.assertEqual(end_col, 21)
 
 
 class TestStaticAttributes(unittest.TestCase):

Lib/test/test_ast/snippets.py

@@ -781,8 +781,8 @@ def exit_raises():
                 return e
 
         for func, expected in [(init_raises, "self.InitRaises()"),
-                               (enter_raises, "self.EnterRaises()"),
-                               (exit_raises, "self.ExitRaises()"),
+                               (enter_raises, "with self.EnterRaises(), self.Dummy() as d:"),
+                               (exit_raises, "with self.ExitRaises(), self.Dummy() as d:"),
                               ]:
             with self.subTest(func):
                 exc = func()

Lib/test/test_compile.py

@@ -35,8 +35,8 @@ module Python
           | AsyncFor(expr target, expr iter, stmt* body, stmt* orelse, string? type_comment)
           | While(expr test, stmt* body, stmt* orelse)
           | If(expr test, stmt* body, stmt* orelse)
-          | With(withitem* items, stmt* body, string? type_comment)
-          | AsyncWith(withitem* items, stmt* body, string? type_comment)
+          | With(withitem* items, stmt* body, string? type_comment, header? head)
+          | AsyncWith(withitem* items, stmt* body, string? type_comment, header? head)
 
           | Match(expr subject, match_case* cases)
 
@@ -149,4 +149,7 @@ module Python
                | ParamSpec(identifier name, expr? default_value)
                | TypeVarTuple(identifier name, expr? default_value)
                attributes (int lineno, int col_offset, int end_lineno, int end_col_offset)
+    
+    header = ()
+             attributes (int lineno, int col_offset, int end_lineno, int end_col_offset)
 }

Lib/test/test_with.py

@@ -1693,4 +1693,4 @@ _PyPegen_concatenate_strings(Parser *p, asdl_expr_seq *strings,
 
     assert(current_pos == n_elements);
     return _PyAST_JoinedStr(values, lineno, col_offset, end_lineno, end_col_offset, p->arena);
-}
+}