From 3513264ceff0e1cfbd9ad6543bc24b5cc08f1acb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Zbigniew=20J=C4=99drzejewski-Szmek?= <zbyszek@in.waw.pl>
Date: Tue, 10 Sep 2019 11:06:29 +0200
Subject: [PATCH 4/4] Make tests pass under python3.8 + astroid 2.2

---
 asttokens/mark_tokens.py  | 27 ++++++++++++++++-----------
 tests/test_mark_tokens.py | 32 +++++++++++++++++++++++---------
 2 files changed, 39 insertions(+), 20 deletions(-)

diff --git a/asttokens/mark_tokens.py b/asttokens/mark_tokens.py
index 624ac6273b..53e1ed6b16 100644
--- a/asttokens/mark_tokens.py
+++ b/asttokens/mark_tokens.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import six
+import sys
 import numbers
 import token
 from . import util
@@ -165,11 +166,14 @@ class MarkTokens(object):
     return (first_token, last_token)
 
   def handle_comp(self, open_brace, node, first_token, last_token):
-    # For list/set/dict comprehensions, we only get the token of the first child, so adjust it to
+    # For list/set/dict comprehensions, in the past, we only got the token of the first child, so adjust it to
     # include the opening brace (the closing brace will be matched automatically).
-    before = self._code.prev_token(first_token)
-    util.expect_token(before, token.OP, open_brace)
-    return (before, last_token)
+    if util.match_token(first_token, token.OP, open_brace):
+      return (first_token, last_token)
+    else:
+      before = self._code.prev_token(first_token)
+      util.expect_token(before, token.OP, open_brace)
+      return (before, last_token)
 
   def visit_listcomp(self, node, first_token, last_token):
     return self.handle_comp('[', node, first_token, last_token)
@@ -224,13 +228,14 @@ class MarkTokens(object):
     return (first_token, self._code.find_token(last_token, token.OP, ']'))
 
   def visit_tuple(self, node, first_token, last_token):
-    # A tuple doesn't include parens; if there is a trailing comma, make it part of the tuple.
-    try:
-      maybe_comma = self._code.next_token(last_token)
-      if util.match_token(maybe_comma, token.OP, ','):
-        last_token = maybe_comma
-    except IndexError:
-      pass
+    if sys.version_info < (3,8):
+      # A tuple doesn't include parens; if there is a trailing comma, make it part of the tuple.
+      try:
+        maybe_comma = self._code.next_token(last_token)
+        if util.match_token(maybe_comma, token.OP, ','):
+          last_token = maybe_comma
+      except IndexError:
+        pass
     return (first_token, last_token)
 
   def visit_str(self, node, first_token, last_token):
diff --git a/tests/test_mark_tokens.py b/tests/test_mark_tokens.py
index 83d35fa432..2dbd7e8b3c 100644
--- a/tests/test_mark_tokens.py
+++ b/tests/test_mark_tokens.py
@@ -98,7 +98,7 @@ b +     # line3
       # All other expressions preserve newlines and comments but are parenthesized.
       'b +     # line3\n  c',
       'b +     # line3\n  c +   # line4\n  d',
-      'a,      # line2\nb +     # line3\n  c +   # line4\n  d',
+      'a,      # line2\nb +     # line3\n  c +   # line4\n  d' if sys.version_info < (3,8) else source,
     })
 
 
@@ -202,8 +202,12 @@ b +     # line3
     # Make sure we don't fail on parsing slices of the form `foo[4:]`.
     source = "(foo.Area_Code, str(foo.Phone)[:3], str(foo.Phone)[3:], foo[:], bar[::, :])"
     m = self.create_mark_checker(source)
-    self.assertEqual(m.view_nodes_at(1, 1),
-                     { "Attribute:foo.Area_Code", "Name:foo", "Tuple:"+source[1:-1] })
+
+    at_1_1 = { "Attribute:foo.Area_Code", "Name:foo" }
+    if sys.version_info < (3,8):
+      at_1_1.add( "Tuple:"+source[1:-1] )
+    self.assertEqual(m.view_nodes_at(1, 1), at_1_1)
+
     self.assertEqual(m.view_nodes_at(1, 16),
                      { "Subscript:str(foo.Phone)[:3]", "Call:str(foo.Phone)", "Name:str"})
     self.assertEqual(m.view_nodes_at(1, 36),
@@ -226,11 +230,14 @@ bar = ('x y z'   # comment2
 """
     m = self.create_mark_checker(source)
     node_name = 'Const' if self.is_astroid_test else 'Constant'
+    joined = self.is_astroid_test or sys.version_info < (3,8)
     self.assertEqual(m.view_nodes_at(2, 6), {
-      node_name + ":'x y z' \\\n'''a b c''' \"u v w\""
+      node_name + ":'x y z'" + (
+        " \\\n'''a b c''' \"u v w\"" if joined else "")
     })
     self.assertEqual(m.view_nodes_at(4, 7), {
-      node_name + ":'x y z'   # comment2\n       'a b c'   # comment3\n       'u v w'"
+      node_name + ":'x y z'" + (
+        "   # comment2\n       'a b c'   # comment3\n       'u v w'" if joined else "")
     })
 
 
@@ -446,16 +453,23 @@ bar = ('x y z'   # comment2
     m.verify_all_nodes(self)
     # The `arguments` node has bogus positions here (and whenever there are no arguments). We
     # don't let that break our test because it's unclear if it matters to anything anyway.
-    self.assertIn('FunctionDef:@deco1\ndef f():\n  pass', m.view_nodes_at(2, 0))
-    self.assertEqual(m.view_nodes_at(2, 1), {'Name:deco1'})
+
+    nodes_2_0 = m.view_nodes_at(2, 0)
+    nodes_2_1 = m.view_nodes_at(2, 1)
+    nodes_5_0 = m.view_nodes_at(5, 0)
+    nodes_5_1 = m.view_nodes_at(5, 1)
+
+    self.assertIn('FunctionDef:@deco1\ndef f():\n  pass', nodes_2_0 or nodes_2_1)
+    self.assertIn('Name:deco1', nodes_2_1)
     if self.is_astroid_test:
       self.assertEqual(m.view_nodes_at(5, 0), {
         'FunctionDef:@deco2(a=1)\ndef g(x):\n  pass',
         'Decorators:@deco2(a=1)'
       })
     else:
-      self.assertEqual(m.view_nodes_at(5, 0), {'FunctionDef:@deco2(a=1)\ndef g(x):\n  pass'})
-    self.assertEqual(m.view_nodes_at(5, 1), {'Name:deco2', 'Call:deco2(a=1)'})
+      self.assertIn('FunctionDef:@deco2(a=1)\ndef g(x):\n  pass', nodes_5_0 or nodes_5_1)
+    self.assertIn('Name:deco2', nodes_5_1)
+    self.assertIn('Call:deco2(a=1)', nodes_5_1)
 
   def test_with(self):
     source = "with foo: pass"