Lexer: Avoid unnecessary EOS after comments

pekkaklarck · yanne · commit 318eacad74a7 · 2019-05-30T19:01:24.000+03:00
diff --git a/src/robot/parsing/lexer/__init__.py b/src/robot/parsing/lexer/__init__.py
@@ -18,7 +18,7 @@
 from .context import TestCaseFileContext
 from .lexers import TestCaseFileLexer
 from .splitter import Splitter
-from .tokens import Token
+from .tokens import EOS, Token
 
 
 class RobotFrameworkLexer(object):
@@ -46,21 +46,19 @@ def get_tokens(self):
         else:
             ignore = {Token.IGNORE}
         for statement in self._handle_old_for(self.statements):
-            name_eos = None
+            name_token = last_token = None
             for token in statement:
                 if token.type in ignore:
                     continue
-                if name_eos:
-                    yield name_eos
-                    name_eos = None
+                if name_token:
+                    yield EOS.from_token(name_token)
+                    name_token = None
                 if token.type == Token.NAME:
-                    name_eos = Token(Token.EOS,
-                                     lineno=token.lineno,
-                                     columnno=token.columnno + len(token.value))
+                    name_token = token
+                last_token = token
                 yield token
-            yield Token(Token.EOS,
-                        lineno=token.lineno,
-                        columnno=token.columnno + len(token.value))
+            if last_token:
+                yield EOS.from_token(last_token)
 
     def _handle_old_for(self, statements):
         end_statement = [Token(Token.END)]
diff --git a/src/robot/parsing/lexer/tokens.py b/src/robot/parsing/lexer/tokens.py
@@ -82,3 +82,13 @@ def __repr__(self):
 
 Token.DATA_TOKENS = [t for t in Token.__dict__
                      if t[0] != '_' and t not in Token.NON_DATA_TOKENS]
+
+
+class EOS(Token):
+
+    def __init__(self, lineno=-1, columnno=-1):
+        Token.__init__(self, Token.EOS, '', lineno, columnno)
+
+    @classmethod
+    def from_token(cls, token):
+        return EOS(token.lineno, token.columnno + len(token.value))
