From 65187930bd9cf8f54233635a99ee7ca04a05eeff Mon Sep 17 00:00:00 2001
From: Claude <github@claude.nl>
Date: Tue, 9 Feb 2016 21:07:18 +0100
Subject: [PATCH] Update tokenizer to adhere to PEP492 magic

---
 jedi/parser/tokenize.py        | 55 ++++++++++++++++++++++++++++++++--
 test/test_parser/test_pgen2.py | 37 +++++++++++++----------
 2 files changed, 74 insertions(+), 18 deletions(-)

diff --git a/jedi/parser/tokenize.py b/jedi/parser/tokenize.py
index 46f9ef76..301d9ff6 100644
--- a/jedi/parser/tokenize.py
+++ b/jedi/parser/tokenize.py
@@ -16,7 +16,7 @@ import re
 from io import StringIO
 from jedi.parser.token import (tok_name, N_TOKENS, ENDMARKER, STRING, NUMBER,
                                NAME, OP, ERRORTOKEN, NEWLINE, INDENT, DEDENT)
-from jedi._compatibility import is_py3
+from jedi._compatibility import is_py3, is_py35
 
 
 cookie_re = re.compile("coding[:=]\s*([-\w.]+)")
@@ -151,7 +151,58 @@ def source_tokens(source):
     """Generate tokens from a the source code (string)."""
     source = source
     readline = StringIO(source).readline
-    return generate_tokens(readline)
+    tokenizer = generate_tokens(readline)
+    if is_py35:
+        tokenizer = fix_async_await_tokenizer(tokenizer)
+    return tokenizer
+
+
+def fix_async_await_tokenizer(tokenizer):
+    # We need to recognise async functions as per
+    # https://www.python.org/dev/peps/pep-0492/#transition-plan
+    # so we need a three item look-ahead
+    from jedi.parser.token import ASYNC, AWAIT
+    assert is_py35
+    try:
+        first = next(tokenizer)
+    except StopIteration:
+        return
+
+    try:
+        second = next(tokenizer)
+    except StopIteration:
+        yield first
+        return
+
+    defs_stack = []
+    indent = 0
+    for item in tokenizer:
+        if first[0] == INDENT:
+            indent += 1
+
+        if first[0] == DEDENT:
+            indent -= 1
+
+        if first[0] == NEWLINE and second[0] != INDENT:
+            while defs_stack and indent <= defs_stack[-1][0]:
+                defs_stack.pop()
+
+        if second[0] == NAME and second[1] == "def" and item[0] == NAME:
+            if first[0] == NAME and first[1] == "async":
+                defs_stack.append((indent, True))
+            else:
+                defs_stack.append((indent, False))
+        in_async_def = defs_stack and defs_stack[-1][1]
+        if in_async_def and first[0] == NAME and first[1] == "async":
+            yield (ASYNC, ) + first[1:]
+        elif in_async_def and first[0] == NAME and first[1] == "await":
+            yield (AWAIT, ) + first[1:]
+        else:
+            yield first
+        first = second
+        second = item
+    yield first
+    yield second
 
 
 def generate_tokens(readline):
diff --git a/test/test_parser/test_pgen2.py b/test/test_parser/test_pgen2.py
index 6d51ae85..576d338f 100644
--- a/test/test_parser/test_pgen2.py
+++ b/test/test_parser/test_pgen2.py
@@ -46,8 +46,8 @@ class GrammarTest(TestCase):
 class TestMatrixMultiplication(GrammarTest):
     @pytest.mark.skipif('sys.version_info[:2] < (3, 5)')
     def test_matrix_multiplication_operator(self):
-        parse("a @ b")
-        parse("a @= b")
+        parse("a @ b", "3.5")
+        parse("a @= b", "3.5")
 
 
 class TestYieldFrom(GrammarTest):
@@ -62,7 +62,7 @@ class TestAsyncAwait(GrammarTest):
     def test_await_expr(self):
         parse("""async def foo():
                              await x
-                      """)
+                      """, "3.5")
 
         parse("""async def foo():
 
@@ -71,46 +71,51 @@ class TestAsyncAwait(GrammarTest):
             def foo(): pass
 
             await x
-        """)
+        """, "3.5")
 
-        parse("""async def foo(): return await a""")
+        parse("""async def foo(): return await a""", "3.5")
 
         parse("""def foo():
             def foo(): pass
             async def foo(): await x
-        """)
+        """, "3.5")
 
-        self.invalid_syntax("await x")
+        self.invalid_syntax("await x", version="3.5")
         self.invalid_syntax("""def foo():
-                                   await x""")
+                                   await x""", version="3.5")
 
         self.invalid_syntax("""def foo():
             def foo(): pass
             async def foo(): pass
             await x
-        """)
+        """, version="3.5")
+
+        self.invalid_syntax("""async def foo():
+            def foo():
+                await x
+        """, version="3.5")
 
     @pytest.mark.skipif('sys.version_info[:2] < (3, 5)')
     def test_async_var(self):
-        parse("""async = 1""")
-        parse("""await = 1""")
-        parse("""def async(): pass""")
+        parse("""async = 1""", "3.5")
+        parse("""await = 1""", "3.5")
+        parse("""def async(): pass""", "3.5")
 
     @pytest.mark.skipif('sys.version_info[:2] < (3, 5)')
     def test_async_with(self):
         parse("""async def foo():
-                             async for a in b: pass""")
+                             async for a in b: pass""", "3.5")
 
         self.invalid_syntax("""def foo():
-                                   async for a in b: pass""")
+                                   async for a in b: pass""", version="3.5")
 
     @pytest.mark.skipif('sys.version_info[:2] < (3, 5)')
     def test_async_for(self):
         parse("""async def foo():
-                             async with a: pass""")
+                             async with a: pass""", "3.5")
 
         self.invalid_syntax("""def foo():
-                                   async with a: pass""")
+                                   async with a: pass""", version="3.5")
 
 
 class TestRaiseChanges(GrammarTest):