From 38050bab6ede05fc96c417f847ad72bc3df1f5fb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Zbigniew=20J=C4=99drzejewski-Szmek?= <zbyszek@in.waw.pl>
Date: Thu, 5 Jul 2018 12:37:33 +0200
Subject: [PATCH] Import typed_ast instead of ast like astroid 2.0 does

astroid 2.0 will attempt to import typed_ast.ast3 and typed_ast.ast27,
and then will fall back to plain ast if that fails. This patch adds
a single import site for ast/typed_ast, and then just loads the result
from other places.

Fixes most of the errors in #16, all similar to:

ERROR: test_calling_lambdas (tests.test_astroid.TestAstroid)
----------------------------------------------------------------------
Traceback (most recent call last):
  File ".../asttokens/tests/test_mark_tokens.py", line 345, in test_calling_lambdas
    m.verify_all_nodes(self)
  File ".../asttokens/tests/tools.py", line 129, in verify_all_nodes
    test_case.assertEqual(to_source(rebuilt_node), to_source(node))
  File ".../asttokens/tests/tools.py", line 59, in to_source
    anode = builder.visit_module(node_copy, '', '', '')
  File ".../asttokens/.tox/py36/lib/python3.6/site-packages/astroid/rebuilder.py", line 152, in visit_module
    newnode.postinit([self.visit(child, newnode) for child in node.body])
  File ".../asttokens/.tox/py36/lib/python3.6/site-packages/astroid/rebuilder.py", line 152, in <listcomp>
    newnode.postinit([self.visit(child, newnode) for child in node.body])
  File ".../asttokens/.tox/py36/lib/python3.6/site-packages/astroid/rebuilder.py", line 164, in visit
    return visit_method(node, parent)
  File ".../asttokens/.tox/py36/lib/python3.6/site-packages/astroid/rebuilder.py", line 268, in visit_assign
    value=self.visit(node.value, newnode),
  File ".../asttokens/.tox/py36/lib/python3.6/site-packages/astroid/rebuilder.py", line 164, in visit
    return visit_method(node, parent)
  File ".../asttokens/.tox/py36/lib/python3.6/site-packages/astroid/rebuilder.py", line 345, in visit_call
    newnode.postinit(self.visit(node.func, newnode),
  File ".../asttokens/.tox/py36/lib/python3.6/site-packages/astroid/rebuilder.py", line 164, in visit
    return visit_method(node, parent)
  File ".../asttokens/.tox/py36/lib/python3.6/site-packages/astroid/rebuilder.py", line 623, in visit_lambda
    self.visit(node.body, newnode))
  File ".../asttokens/.tox/py36/lib/python3.6/site-packages/astroid/rebuilder.py", line 164, in visit
    return visit_method(node, parent)
  File ".../asttokens/.tox/py36/lib/python3.6/site-packages/astroid/rebuilder.py", line 296, in visit_binop
    newnode = nodes.BinOp(self._bin_op_classes[type(node.op)],
KeyError: <class '_ast.Add'>

self._bin_op_classes contains '_ast3.Add', i.e. typed_ast._ast3.Add instead.
---
 asttokens/asttokens.py  |  3 +--
 asttokens/util.py       | 13 ++++++++++++-
 tests/test_asttokens.py |  2 +-
 tests/test_util.py      |  2 +-
 tests/tools.py          |  3 +--
 5 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/asttokens/asttokens.py b/asttokens/asttokens.py
index c0126a6b78..4659d3327e 100644
--- a/asttokens/asttokens.py
+++ b/asttokens/asttokens.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import ast
 import bisect
 import token
 import tokenize
@@ -20,7 +19,7 @@ import io
 import six
 from six.moves import xrange      # pylint: disable=redefined-builtin
 from .line_numbers import LineNumbers
-from .util import Token, match_token
+from .util import ast, Token, match_token
 from .mark_tokens import MarkTokens
 
 class ASTTokens(object):
diff --git a/asttokens/util.py b/asttokens/util.py
index 4dd2f27983..29d1e7e0cb 100644
--- a/asttokens/util.py
+++ b/asttokens/util.py
@@ -12,11 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import ast
 import collections
 import token
+import sys
 from six import iteritems
 
+import astroid
+if astroid.__version__ < '2':
+  import ast
+else:
+  try:
+    if sys.version_info.major >= 3:
+      import typed_ast.ast3 as ast
+    else:
+      import typed_ast.ast27 as ast
+  except ImportError:
+    import ast
 
 def token_repr(tok_type, string):
   """Returns a human-friendly representation of a token with the given type and string."""
diff --git a/tests/test_asttokens.py b/tests/test_asttokens.py
index 213eca288b..67b0de8597 100644
--- a/tests/test_asttokens.py
+++ b/tests/test_asttokens.py
@@ -1,11 +1,11 @@
 # -*- coding: UTF-8 -*-
 from __future__ import unicode_literals, print_function
-import ast
 import six
 import token
 import tokenize
 import unittest
 from .context import asttokens
+ast = asttokens.util.ast
 from . import tools
 
 class TestASTTokens(unittest.TestCase):
diff --git a/tests/test_util.py b/tests/test_util.py
index 2dd09350fb..ceafc1d182 100644
--- a/tests/test_util.py
+++ b/tests/test_util.py
@@ -1,9 +1,9 @@
 # -*- coding: UTF-8 -*-
 from __future__ import unicode_literals, print_function
-import ast
 import astroid
 import unittest
 from .context import asttokens
+ast = asttokens.util.ast
 
 class TestUtil(unittest.TestCase):
 
diff --git a/tests/tools.py b/tests/tools.py
index b98f061c89..b8db8fa9ae 100644
--- a/tests/tools.py
+++ b/tests/tools.py
@@ -1,5 +1,4 @@
 from __future__ import unicode_literals, print_function
-import ast
 import astroid
 import asttokens
 import copy
@@ -8,7 +7,7 @@ import os
 import re
 import sys
 from asttokens import util
-
+ast = util.ast
 
 def get_fixture_path(*path_parts):
   python_dir = 'python%s' % sys.version_info[0]