Fix col_offset attribute for nodes involving with on PyPy (#1520)

Co-authored-by: Marc Mueller <30130371+cdce8p@users.noreply.github.com>

Author: DanielNoord

ChangeLog

@@ -13,6 +13,7 @@ What's New in astroid 2.11.4?
 =============================
 Release date: TBA
 
+* Fix ``col_offset`` attribute for nodes involving ``with`` on ``PyPy``.
 
 
 What's New in astroid 2.11.3?

astroid/rebuilder.py

@@ -6,13 +6,13 @@
 order to get a single Astroid representation
 """
 
+import ast
 import sys
 import token
 import tokenize
 from io import StringIO
 from tokenize import TokenInfo, generate_tokens
 from typing import (
-    TYPE_CHECKING,
     Callable,
     Dict,
     Generator,
@@ -39,9 +39,6 @@
 else:
     from typing_extensions import Final
 
-if TYPE_CHECKING:
-    import ast
-
 
 REDIRECT: Final[Dict[str, str]] = {
     "arguments": "Arguments",
@@ -1185,9 +1182,14 @@ def _visit_for(
         self, cls: Type[T_For], node: Union["ast.For", "ast.AsyncFor"], parent: NodeNG
     ) -> T_For:
         """visit a For node by returning a fresh instance of it"""
+        col_offset = node.col_offset
+        if IS_PYPY and not PY39_PLUS and isinstance(node, ast.AsyncFor) and self._data:
+            # pylint: disable-next=unsubscriptable-object
+            col_offset = self._data[node.lineno - 1].index("async")
+
         newnode = cls(
             lineno=node.lineno,
-            col_offset=node.col_offset,
+            col_offset=col_offset,
             # end_lineno and end_col_offset added in 3.8
             end_lineno=getattr(node, "end_lineno", None),
             end_col_offset=getattr(node, "end_col_offset", None),
@@ -1292,6 +1294,10 @@ def _visit_functiondef(
             position=self._get_position_info(node, newnode),
             doc_node=self.visit(doc_ast_node, newnode),
         )
+        if IS_PYPY and PY36 and newnode.position:
+            # PyPy: col_offset in Python 3.6 doesn't include 'async',
+            # use position.col_offset instead.
+            newnode.col_offset = newnode.position.col_offset
         self._fix_doc_node_position(newnode)
         self._global_names.pop()
         return newnode
@@ -1903,9 +1909,14 @@ def _visit_with(
         node: Union["ast.With", "ast.AsyncWith"],
         parent: NodeNG,
     ) -> T_With:
+        col_offset = node.col_offset
+        if IS_PYPY and not PY39_PLUS and isinstance(node, ast.AsyncWith) and self._data:
+            # pylint: disable-next=unsubscriptable-object
+            col_offset = self._data[node.lineno - 1].index("async")
+
         newnode = cls(
             lineno=node.lineno,
-            col_offset=node.col_offset,
+            col_offset=col_offset,
             # end_lineno and end_col_offset added in 3.8
             end_lineno=getattr(node, "end_lineno", None),
             end_col_offset=getattr(node, "end_col_offset", None),

tests/unittest_nodes.py

@@ -1074,20 +1074,52 @@ def test(a): return a
 
 class Python35AsyncTest(unittest.TestCase):
     def test_async_await_keywords(self) -> None:
-        async_def, async_for, async_with, await_node = builder.extract_node(
+        (
+            async_def,
+            async_for,
+            async_with,
+            async_for2,
+            async_with2,
+            await_node,
+        ) = builder.extract_node(
             """
         async def func(): #@
             async for i in range(10): #@
                 f = __(await i)
             async with test(): #@
                 pass
+            async for i \
+                    in range(10):  #@
+                pass
+            async with test(), \
+                    test2():  #@
+                pass
         """
         )
-        self.assertIsInstance(async_def, nodes.AsyncFunctionDef)
-        self.assertIsInstance(async_for, nodes.AsyncFor)
-        self.assertIsInstance(async_with, nodes.AsyncWith)
-        self.assertIsInstance(await_node, nodes.Await)
-        self.assertIsInstance(await_node.value, nodes.Name)
+        assert isinstance(async_def, nodes.AsyncFunctionDef)
+        assert async_def.lineno == 2
+        assert async_def.col_offset == 0
+
+        assert isinstance(async_for, nodes.AsyncFor)
+        assert async_for.lineno == 3
+        assert async_for.col_offset == 4
+
+        assert isinstance(async_with, nodes.AsyncWith)
+        assert async_with.lineno == 5
+        assert async_with.col_offset == 4
+
+        assert isinstance(async_for2, nodes.AsyncFor)
+        assert async_for2.lineno == 7
+        assert async_for2.col_offset == 4
+
+        assert isinstance(async_with2, nodes.AsyncWith)
+        assert async_with2.lineno == 9
+        assert async_with2.col_offset == 4
+
+        assert isinstance(await_node, nodes.Await)
+        assert isinstance(await_node.value, nodes.Name)
+        assert await_node.lineno == 4
+        assert await_node.col_offset == 15
 
     def _test_await_async_as_string(self, code: str) -> None:
         ast_node = parse(code)