From a588a70a7a691af7bef235f38835777c414e91b0 Mon Sep 17 00:00:00 2001
From: Erez Shinan <erezshin+git@gmail.com>
Date: Sat, 27 May 2017 01:55:49 +0300
Subject: [PATCH] Added the experimental "propagate_positions" feature (only
 for standard lexer for now).

---
 lark/lark.py               |  4 +++-
 lark/lexer.py              | 16 +++++++++++++---
 lark/parse_tree_builder.py | 27 ++++++++++++++++++++++++++-
 lark/utils.py              | 20 ++++++++++++++++++++
 4 files changed, 62 insertions(+), 5 deletions(-)

diff --git a/lark/lark.py b/lark/lark.py
index b839650..3d8cbcd 100644
--- a/lark/lark.py
+++ b/lark/lark.py
@@ -39,6 +39,7 @@ class LarkOptions(object):
         postlex - Lexer post-processing (Default: None)
         start - The start symbol (Default: start)
         profile - Measure run-time usage in Lark. Read results from the profiler proprety (Default: False)
+        propagate_positions - Experimental. Don't use yet.
     """
     __doc__ += OPTIONS_DOC
     def __init__(self, options_dict):
@@ -55,6 +56,7 @@ class LarkOptions(object):
         self.start = o.pop('start', 'start')
         self.profile = o.pop('profile', False)
         self.ambiguity = o.pop('ambiguity', 'auto')
+        self.propagate_positions = o.pop('propagate_positions', False)
 
         assert self.parser in ('earley', 'lalr', None)
 
@@ -160,7 +162,7 @@ class Lark:
 
     def _build_parser(self):
         self.parser_class = get_frontend(self.options.parser, self.options.lexer)
-        self.parse_tree_builder = ParseTreeBuilder(self.options.tree_class)
+        self.parse_tree_builder = ParseTreeBuilder(self.options.tree_class, self.options.propagate_positions)
         rules, callback = self.parse_tree_builder.create_tree_builder(self.rules, self.options.transformer)
         if self.profiler:
             for f in dir(callback):
diff --git a/lark/lexer.py b/lark/lexer.py
index 053ce32..63e306f 100644
--- a/lark/lexer.py
+++ b/lark/lexer.py
@@ -155,17 +155,27 @@ class Lexer(object):
                 if m:
                     value = m.group(0)
                     type_ = type_from_index[m.lastindex]
-                    if type_ not in ignore_types:
+                    to_yield = type_ not in ignore_types
+
+                    if to_yield:
                         t = Token(type_, value, lex_pos, line, lex_pos - col_start_pos)
                         if t.type in self.callback:
                             t = self.callback[t.type](t)
-                        yield t
 
+                    end_col = t.column + len(value)
                     if type_ in newline_types:
                         newlines = value.count(self.newline_char)
                         if newlines:
                             line += newlines
-                            col_start_pos = lex_pos + value.rindex(self.newline_char)
+                            last_newline_index = value.rindex(self.newline_char) + 1
+                            col_start_pos = lex_pos + last_newline_index
+                            end_col = len(value) - last_newline_index
+
+                    if to_yield:
+                        t.end_line = line
+                        t.end_col = end_col
+                        yield t
+
                     lex_pos += len(value)
                     break
             else:
diff --git a/lark/parse_tree_builder.py b/lark/parse_tree_builder.py
index 547deab..b3bc522 100644
--- a/lark/parse_tree_builder.py
+++ b/lark/parse_tree_builder.py
@@ -1,4 +1,5 @@
 from .common import is_terminal, GrammarError
+from .utils import suppress
 from .lexer import Token
 
 class Callback(object):
@@ -42,10 +43,31 @@ def create_rule_handler(expansion, usermethod, keep_all_tokens, filter_out):
     # else, if no filtering required..
     return usermethod
 
+def propagate_positions_wrapper(f):
+    def _f(args):
+        res = f(args)
+
+        if args:
+            for a in args:
+                with suppress(AttributeError):
+                    res.line = a.line
+                    res.column = a.column
+                break
+
+            for a in reversed(args):
+                with suppress(AttributeError):
+                    res.end_line = a.end_line
+                    res.end_col = a.end_col
+                break
+
+        return res
+
+    return _f
 
 class ParseTreeBuilder:
-    def __init__(self, tree_class):
+    def __init__(self, tree_class, propagate_positions=False):
         self.tree_class = tree_class
+        self.propagate_positions = propagate_positions
 
     def _create_tree_builder_function(self, name):
         tree_class = self.tree_class
@@ -92,6 +114,9 @@ class ParseTreeBuilder:
 
                 alias_handler = create_rule_handler(expansion, f, keep_all_tokens, filter_out)
 
+                if self.propagate_positions:
+                    alias_handler = propagate_positions_wrapper(alias_handler)
+
                 callback_name = 'autoalias_%s_%s' % (_origin, '_'.join(expansion))
                 if hasattr(callback, callback_name):
                     raise GrammarError("Rule expansion '%s' already exists in rule %s" % (' '.join(expansion), origin))
diff --git a/lark/utils.py b/lark/utils.py
index 91b8a24..d984400 100644
--- a/lark/utils.py
+++ b/lark/utils.py
@@ -1,6 +1,7 @@
 import functools
 import types
 from collections import deque
+from contextlib import contextmanager
 
 class fzset(frozenset):
     def __repr__(self):
@@ -88,5 +89,24 @@ except NameError:
             return -1
 
 
+try:
+    from contextlib import suppress     # Python 3
+except ImportError:
+    @contextmanager
+    def suppress(*excs):
+        '''Catch and dismiss the provided exception
+
+        >>> x = 'hello'
+        >>> with suppress(IndexError):
+        ...     x = x[10]
+        >>> x
+        'hello'
+        '''
+        try:
+            yield
+        except excs:
+            pass
+
+