Support [g]ast.get_source_segment

Based on py3.8 implementation

Author: serge-sans-paille

gast/gast.py

@@ -461,6 +461,77 @@ def _fix(node, lineno, col_offset, end_lineno, end_col_offset):
     return node
 
 
+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 increment_lineno(node, n=1):
     """
     Increment the line number and end line number of each node in the tree

tests/test_api.py

@@ -74,6 +74,21 @@ def test_increment_lineno(self):
         gast.increment_lineno(tree)
         self.assertEqual(tree.lineno, 2)
 
+    def test_get_source_segment(self):
+        code = 'x + 1'
+        tree = gast.parse(code)
+        source = gast.get_source_segment(code, tree.body[0].value.left)
+        self.assertEqual(source, 'x')
+
+    def test_get_source_segment_padded(self):
+        code = 'if 1:\n if 2:\n  3'
+        tree = gast.parse(code)
+        if_tree = tree.body[0].body[0]
+        source_nopadding = gast.get_source_segment(code, if_tree, padded=False)
+        self.assertEqual(source_nopadding, 'if 2:\n  3')
+        source_padding = gast.get_source_segment(code, if_tree, padded=True)
+        self.assertEqual(source_padding, ' if 2:\n  3')
+
     def test_get_docstring_function(self):
         code = 'def foo(): "foo"'
         tree = gast.parse(code)