end_pos issues, fixes #150

jedi/common.py

@@ -64,6 +64,7 @@ class PushBackIterator(object):
     def __init__(self, iterator):
         self.pushes = []
         self.iterator = iterator
+        self.current = None
 
     def push_back(self, value):
         self.pushes.append(value)
@@ -77,9 +78,10 @@ class PushBackIterator(object):
 
     def __next__(self):
         if self.pushes:
-            return self.pushes.pop()
+            self.current = self.pushes.pop()
         else:
-            return next(self.iterator)
+            self.current = next(self.iterator)
+        return self.current
 
 
 class NoErrorTokenizer(object):

jedi/parsing_representation.py

@@ -884,12 +884,9 @@ class Statement(Simple):
                 # always dictionaries and not sets.
                 arr.type = Array.DICT
 
-            k, v = arr.keys, arr.values
-            latest = (v[-1] if v else k[-1] if k else None)
-            end_pos = latest.end_pos if latest is not None \
-                                     else (start_pos[0], start_pos[1] + 1)
-            arr.end_pos = end_pos[0], end_pos[1] + (len(break_tok) if break_tok
-                                                    else 0)
+            c = token_iterator.current[1]
+            arr.end_pos = c.end_pos if isinstance(c, Simple) \
+                                        else (c[2][0], c[2][1] + len(c[1]))
             return arr, break_tok
 
         def parse_stmt(token_iterator, maybe_dict=False, added_breaks=(),

test/test_regression.py

@@ -15,7 +15,7 @@ from .base import TestBase, unittest, cwd_at
 
 import jedi
 from jedi._compatibility import utf8, unicode
-from jedi import api
+from jedi import api, parsing
 api_classes = api.api_classes
 
 import pytest
@@ -378,6 +378,14 @@ class TestRegression(TestBase):
         self.assertEqual([d.description for d in defs],
                          ['def f', 'class C'])
 
+    def test_end_pos(self):
+        # jedi issue #150
+        s = "x()\nx( )\nx(  )\nx (  )"
+        parser = parsing.Parser(s)
+        for i, s in enumerate(parser.scope.statements, 3):
+            for c in s.get_commands():
+                self.assertEqual(c.execution.end_pos[1], i)
+
 
 class TestDocstring(TestBase):