bpo-33416: Add end positions to Python AST

Doc/library/ast.rst

@@ -61,13 +61,21 @@ Node classes
 
    .. attribute:: lineno
                   col_offset
+                  end_lineno
+                  end_col_offset
 
       Instances of :class:`ast.expr` and :class:`ast.stmt` subclasses have
-      :attr:`lineno` and :attr:`col_offset` attributes.  The :attr:`lineno` is
-      the line number of source text (1-indexed so the first line is line 1) and
-      the :attr:`col_offset` is the UTF-8 byte offset of the first token that
-      generated the node.  The UTF-8 offset is recorded because the parser uses
-      UTF-8 internally.
+      :attr:`lineno`, :attr:`col_offset`, :attr:`lineno`, and :attr:`col_offset`
+      attributes.  The :attr:`lineno` and :attr:`end_lineno` are the first and
+      last line numbers of source text span (1-indexed so the first line is line 1)
+      and the :attr:`col_offset` and :attr:`end_col_offset` are the corresponding
+      UTF-8 byte offsets of the first and last tokens that generated the node.
+      The UTF-8 offset is recorded because the parser uses UTF-8 internally.
+
+      Note that the end positions are not required by the compiler and are
+      therefore optional. The end offset is *after* the last symbol, for example
+      one can get the source segment of a one-line expression node using
+      ``source_line[node.col_offset : node.end_col_offset]``.
 
    The constructor of a class :class:`ast.T` parses its arguments as follows:
 
@@ -162,6 +170,18 @@ and classes for traversing abstract syntax trees:
       :class:`AsyncFunctionDef` is now supported.
 
 
+.. function:: get_source_segment(source, node, *, padded=False)
+
+   Get source code segment of the *source* that generated *node*.
+   If some location information (:attr:`lineno`, :attr:`end_lineno`,
+   :attr:`col_offset`, or :attr:`end_col_offset`) is missing, return ``None``.
+
+   If *padded* is ``True``, the first line of a multi-line statement will
+   be padded with spaces to match its original position.
+
+   .. versionadded:: 3.8
+
+
 .. function:: fix_missing_locations(node)
 
    When you compile a node tree with :func:`compile`, the compiler expects
@@ -173,14 +193,16 @@ and classes for traversing abstract syntax trees:
 
 .. function:: increment_lineno(node, n=1)
 
-   Increment the line number of each node in the tree starting at *node* by *n*.
-   This is useful to "move code" to a different location in a file.
+   Increment the line number and end line number of each node in the tree
+   starting at *node* by *n*. This is useful to "move code" to a different
+   location in a file.
 
 
 .. function:: copy_location(new_node, old_node)
 
-   Copy source location (:attr:`lineno` and :attr:`col_offset`) from *old_node*
-   to *new_node* if possible, and return *new_node*.
+   Copy source location (:attr:`lineno`, :attr:`col_offset`, :attr:`end_lineno`,
+   and :attr:`end_col_offset`) from *old_node* to *new_node* if possible,
+   and return *new_node*.
 
 
 .. function:: iter_fields(node)

Include/node.h

@@ -14,11 +14,14 @@ typedef struct _node {
     int                 n_col_offset;
     int                 n_nchildren;
     struct _node        *n_child;
+    int                 n_end_lineno;
+    int                 n_end_col_offset;
 } node;
 
 PyAPI_FUNC(node *) PyNode_New(int type);
 PyAPI_FUNC(int) PyNode_AddChild(node *n, int type,
-                                      char *str, int lineno, int col_offset);
+                                char *str, int lineno, int col_offset,
+                                int end_lineno, int end_col_offset);
 PyAPI_FUNC(void) PyNode_Free(node *n);
 #ifndef Py_LIMITED_API
 PyAPI_FUNC(Py_ssize_t) _PyNode_SizeOf(node *n);
@@ -37,6 +40,7 @@ PyAPI_FUNC(Py_ssize_t) _PyNode_SizeOf(node *n);
 #define REQ(n, type) assert(TYPE(n) == (type))
 
 PyAPI_FUNC(void) PyNode_ListTree(node *);
+void _PyNode_FinalizeEndPos(node *n);  // helper also used in parsetok.c
 
 #ifdef __cplusplus
 }

Lib/ast.py

@@ -115,10 +115,10 @@ def _format(node):
 
 def copy_location(new_node, old_node):
     """
-    Copy source location (`lineno` and `col_offset` attributes) from
-    *old_node* to *new_node* if possible, and return *new_node*.
+    Copy source location (`lineno`, `col_offset`, `end_lineno`, and `end_col_offset`
+    attributes) from *old_node* to *new_node* if possible, and return *new_node*.
     """
-    for attr in 'lineno', 'col_offset':
+    for attr in 'lineno', 'col_offset', 'end_lineno', 'end_col_offset':
         if attr in old_node._attributes and attr in new_node._attributes \
            and hasattr(old_node, attr):
             setattr(new_node, attr, getattr(old_node, attr))
@@ -133,31 +133,44 @@ def fix_missing_locations(node):
     recursively where not already set, by setting them to the values of the
     parent node.  It works recursively starting at *node*.
     """
-    def _fix(node, lineno, col_offset):
+    def _fix(node, lineno, col_offset, end_lineno, end_col_offset):
         if 'lineno' in node._attributes:
             if not hasattr(node, 'lineno'):
                 node.lineno = lineno
             else:
                 lineno = node.lineno
+        if 'end_lineno' in node._attributes:
+            if not hasattr(node, 'end_lineno'):
+                node.end_lineno = end_lineno
+            else:
+                end_lineno = node.end_lineno
         if 'col_offset' in node._attributes:
             if not hasattr(node, 'col_offset'):
                 node.col_offset = col_offset
             else:
                 col_offset = node.col_offset
+        if 'end_col_offset' in node._attributes:
+            if not hasattr(node, 'end_col_offset'):
+                node.end_col_offset = end_col_offset
+            else:
+                end_col_offset = node.end_col_offset
         for child in iter_child_nodes(node):
-            _fix(child, lineno, col_offset)
-    _fix(node, 1, 0)
+            _fix(child, lineno, col_offset, end_lineno, end_col_offset)
+    _fix(node, 1, 0, 1, 0)
     return node
 
 
 def increment_lineno(node, n=1):
     """
-    Increment the line number of each node in the tree starting at *node* by *n*.
-    This is useful to "move code" to a different location in a file.
+    Increment the line number and end line number of each node in the tree
+    starting at *node* by *n*. This is useful to "move code" to a different
+    location in a file.
     """
     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
     return node
 
 
@@ -213,6 +226,77 @@ def get_docstring(node, clean=True):
     return text
 
 
+def _splitlines_no_ff(source):
+    """Split a string into lines ignoring form feed and other chars.
+
+    This mimics how the Python parser splits source code.
+    """
+    idx = 0
+    lines = []
+    next_line = ''
+    while idx < len(source):
+        c = source[idx]
+        next_line += c
+        idx += 1
+        # Keep \r\n together
+        if c == '\r' and idx < len(source) and source[idx] == '\n':
+            next_line += '\n'
+            idx += 1
+        if c in '\r\n':
+            lines.append(next_line)
+            next_line = ''
+
+    if next_line:
+        lines.append(next_line)
+    return lines
+
+
+def _pad_whitespace(source):
+    """Replace all chars except '\f\t' in a line with spaces."""
+    result = ''
+    for c in source:
+        if c in '\f\t':
+            result += c
+        else:
+            result += ' '
+    return result
+
+
+def get_source_segment(source, node, *, padded=False):
+    """Get source code segment of the *source* that generated *node*.
+
+    If some location information (`lineno`, `end_lineno`, `col_offset`,
+    or `end_col_offset`) is missing, return None.
+
+    If *padded* is `True`, the first line of a multi-line statement will
+    be padded with spaces to match its original position.
+    """
+    try:
+        lineno = node.lineno - 1
+        end_lineno = node.end_lineno - 1
+        col_offset = node.col_offset
+        end_col_offset = node.end_col_offset
+    except AttributeError:
+        return None
+
+    lines = _splitlines_no_ff(source)
+    if end_lineno == lineno:
+        return lines[lineno].encode()[col_offset:end_col_offset].decode()
+
+    if padded:
+        padding = _pad_whitespace(lines[lineno].encode()[:col_offset].decode())
+    else:
+        padding = ''
+
+    first = padding + lines[lineno].encode()[col_offset:].decode()
+    last = lines[end_lineno].encode()[:end_col_offset].decode()
+    lines = lines[lineno+1:end_lineno]
+
+    lines.insert(0, first)
+    lines.append(last)
+    return ''.join(lines)
+
+
 def walk(node):
     """
     Recursively yield all descendant nodes in the tree starting at *node*

Lib/test/test_asdl_parser.py

@@ -62,14 +62,16 @@ def test_product(self):
 
     def test_attributes(self):
         stmt = self.types['stmt']
-        self.assertEqual(len(stmt.attributes), 2)
+        self.assertEqual(len(stmt.attributes), 4)
         self.assertEqual(str(stmt.attributes[0]), 'Field(int, lineno)')
         self.assertEqual(str(stmt.attributes[1]), 'Field(int, col_offset)')
+        self.assertEqual(str(stmt.attributes[2]), 'Field(int, end_lineno, opt=True)')
+        self.assertEqual(str(stmt.attributes[3]), 'Field(int, end_col_offset, opt=True)')
 
     def test_constructor_fields(self):
         ehandler = self.types['excepthandler']
         self.assertEqual(len(ehandler.types), 1)
-        self.assertEqual(len(ehandler.attributes), 2)
+        self.assertEqual(len(ehandler.attributes), 4)
 
         cons = ehandler.types[0]
         self.assertIsInstance(cons, self.asdl.Constructor)