From 35c18ee0823f85102efcfac7c21499c8d6a2244f Mon Sep 17 00:00:00 2001
From: Samuel Colvin <s@muelcolvin.com>
Date: Sat, 30 Sep 2017 13:57:12 +0100
Subject: [PATCH] allow async await debug, fix #3

---
 .gitignore                |  1 +
 Makefile                  |  6 +++---
 dev-requirements.txt      |  1 +
 devtools/debug.py         | 22 +++++++++++++++-------
 tests/test_expr_render.py | 17 +++++++++++++++++
 5 files changed, 37 insertions(+), 10 deletions(-)

diff --git a/.gitignore b/.gitignore
index c2d9678..3bfcf33 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,4 +1,5 @@
 .idea/
+env/
 env35/
 env36/
 *.py[cod]
diff --git a/Makefile b/Makefile
index 1d0a64d..61f255f 100644
--- a/Makefile
+++ b/Makefile
@@ -1,12 +1,12 @@
 .DEFAULT_GOAL := all
 
-.PHONY: install
+.PHONY: install-dev
 install-dev:
 	pip install -U setuptools pip
 	pip install -r dev-requirements.txt
-	pip install -U .
+	pip install -U -e .[pygments]
 
-.PHONY: install
+.PHONY: install-test
 install-test:
 	pip install -U setuptools pip
 	pip install -r tests/requirements.txt
diff --git a/dev-requirements.txt b/dev-requirements.txt
index 91e04ab..1cb9ead 100644
--- a/dev-requirements.txt
+++ b/dev-requirements.txt
@@ -1,4 +1,5 @@
 #-r docs/requirements.txt
 -r tests/requirements.txt
 
+astpretty  # pyup: ignore
 ipython  # pyup: ignore
diff --git a/devtools/debug.py b/devtools/debug.py
index 10a5b6a..7c26941 100644
--- a/devtools/debug.py
+++ b/devtools/debug.py
@@ -5,7 +5,7 @@
 import re
 import warnings
 from pathlib import Path
-from textwrap import dedent
+from textwrap import dedent, indent
 from typing import Generator, List, Optional, Tuple, Type
 
 from .ansi import isatty, sformat
@@ -222,7 +222,7 @@ def _parse_code(self, call_frame, func_regex, filename) -> Tuple[Optional[ast.AS
         func_ast = None
         tail_index = call_frame.index
         try:
-            func_ast = ast.parse(code, filename=filename).body[0].value
+            func_ast = self._wrap_parse(code, filename)
         except SyntaxError as e1:
             # if the trailing bracket(s) of the function is/are on a new line eg.
             # debug(
@@ -233,7 +233,7 @@ def _parse_code(self, call_frame, func_regex, filename) -> Tuple[Optional[ast.AS
                 extra_lines = call_frame.code_context[tail_index + 1:tail_index + extra]
                 code = dedent(''.join(call_lines + extra_lines))
                 try:
-                    func_ast = ast.parse(code).body[0].value
+                    func_ast = self._wrap_parse(code, filename)
                 except SyntaxError:
                     pass
                 else:
@@ -253,17 +253,25 @@ def _parse_code(self, call_frame, func_regex, filename) -> Tuple[Optional[ast.AS
         code_lines[-1] = code_lines[-1][:-1]
         return func_ast, code_lines, lineno
 
-    @classmethod
-    def _get_offsets(cls, func_ast):
+    @staticmethod
+    def _wrap_parse(code, filename):
+        """
+        async wrapper is required to avoid await calls raising a SyntaxError
+        """
+        code = 'async def wrapper():\n' + indent(code, ' ')
+        return ast.parse(code, filename=filename).body[0].body[0].value
+
+    @staticmethod
+    def _get_offsets(func_ast):
         for arg in func_ast.args:
-            start_line, start_col = arg.lineno - 1, arg.col_offset
+            start_line, start_col = arg.lineno - 2, arg.col_offset - 1
 
             # horrible hack for http://bugs.python.org/issue31241
             if isinstance(arg, (ast.ListComp, ast.GeneratorExp)):
                 start_col -= 1
             yield start_line, start_col
         for kw in func_ast.keywords:
-            yield kw.value.lineno - 1, kw.value.col_offset - len(kw.arg) - 1
+            yield kw.value.lineno - 2, kw.value.col_offset - len(kw.arg) - 2
 
     def _warn(self, msg, category: Type[Warning]=RuntimeWarning):
         if self._show_warnings:
diff --git a/tests/test_expr_render.py b/tests/test_expr_render.py
index c7e85f0..4ce8e4e 100644
--- a/tests/test_expr_render.py
+++ b/tests/test_expr_render.py
@@ -1,3 +1,4 @@
+import asyncio
 import re
 import sys
 
@@ -199,3 +200,19 @@ def test_no_syntax_warning():
         )
         assert 'test_no_syntax_warning\n    1 (int)' in str(v)
     assert len(warning_checker) == 0
+
+
+def test_await():
+    async def foo():
+        return 1
+
+    async def bar():
+        return debug.format(await foo())
+
+    loop = asyncio.get_event_loop()
+    v = loop.run_until_complete(bar())
+    s = re.sub(':\d{2,}', ':<line no>', str(v))
+    assert (
+        'tests/test_expr_render.py:<line no> bar\n'
+        '    1 (int)'
+    ) == s