bpo-40726: handle uninitalized end_lineno on ast.increment_lineno (GH-20312)
diff --git a/Lib/ast.py b/Lib/ast.py
index 6a5b39e..65ebd01 100644
--- a/Lib/ast.py
+++ b/Lib/ast.py
@@ -180,7 +180,11 @@
for attr in 'lineno', 'col_offset', 'end_lineno', 'end_col_offset':
if attr in old_node._attributes and attr in new_node._attributes:
value = getattr(old_node, attr, None)
- if value is not None:
+ # end_lineno and end_col_offset are optional attributes, and they
+ # should be copied whether the value is None or not.
+ if value is not None or (
+ hasattr(old_node, attr) and attr.startswith("end_")
+ ):
setattr(new_node, attr, value)
return new_node
@@ -229,8 +233,11 @@
for child in walk(node):
if 'lineno' in child._attributes:
child.lineno = getattr(child, 'lineno', 0) + n
- if 'end_lineno' in child._attributes:
- child.end_lineno = getattr(child, 'end_lineno', 0) + n
+ if (
+ "end_lineno" in child._attributes
+ and (end_lineno := getattr(child, "end_lineno", 0)) is not None
+ ):
+ child.end_lineno = end_lineno + n
return node
diff --git a/Lib/test/test_ast.py b/Lib/test/test_ast.py
index 78e4a56..f5aef61 100644
--- a/Lib/test/test_ast.py
+++ b/Lib/test/test_ast.py
@@ -812,6 +812,12 @@
'lineno=1, col_offset=4, end_lineno=1, end_col_offset=5), lineno=1, '
'col_offset=0, end_lineno=1, end_col_offset=5))'
)
+ src = ast.Call(col_offset=1, lineno=1, end_lineno=1, end_col_offset=1)
+ new = ast.copy_location(src, ast.Call(col_offset=None, lineno=None))
+ self.assertIsNone(new.end_lineno)
+ self.assertIsNone(new.end_col_offset)
+ self.assertEqual(new.lineno, 1)
+ self.assertEqual(new.col_offset, 1)
def test_fix_missing_locations(self):
src = ast.parse('write("spam")')
@@ -851,6 +857,11 @@
'lineno=4, col_offset=4, end_lineno=4, end_col_offset=5), lineno=4, '
'col_offset=0, end_lineno=4, end_col_offset=5))'
)
+ src = ast.Call(
+ func=ast.Name("test", ast.Load()), args=[], keywords=[], lineno=1
+ )
+ self.assertEqual(ast.increment_lineno(src).lineno, 2)
+ self.assertIsNone(ast.increment_lineno(src).end_lineno)
def test_iter_fields(self):
node = ast.parse('foo()', mode='eval')