From fef682a53b1fd22a678b38b0ba6e82c53506f55b Mon Sep 17 00:00:00 2001
From: hauntsaninja <>
Date: Sun, 11 Apr 2021 15:38:32 -0700
Subject: [PATCH 1/2] Fix assertion rewriting on Python 3.10

Fixes https://github.com/pytest-dev/pytest/issues/8539

This seems to have been the result of https://bugs.python.org/issue43798
---
 AUTHORS                          |  1 +
 changelog/8539.bugfix.rst        |  1 +
 src/_pytest/assertion/rewrite.py | 18 ++++++++++++++----
 3 files changed, 16 insertions(+), 4 deletions(-)
 create mode 100644 changelog/8539.bugfix.rst

diff --git a/AUTHORS b/AUTHORS
index 80fce539294..593e0e91a55 100644
--- a/AUTHORS
+++ b/AUTHORS
@@ -222,6 +222,7 @@ Samuele Pedroni
 Sankt Petersbug
 Segev Finer
 Serhii Mozghovyi
+Shantanu Jain
 Simon Gomizelj
 Skylar Downes
 Srinivas Reddy Thatiparthy
diff --git a/changelog/8539.bugfix.rst b/changelog/8539.bugfix.rst
new file mode 100644
index 00000000000..a2098610e29
--- /dev/null
+++ b/changelog/8539.bugfix.rst
@@ -0,0 +1 @@
+Fixed assertion rewriting on Python 3.10.
diff --git a/src/_pytest/assertion/rewrite.py b/src/_pytest/assertion/rewrite.py
index 1c6161b212d..3cc17ad4526 100644
--- a/src/_pytest/assertion/rewrite.py
+++ b/src/_pytest/assertion/rewrite.py
@@ -652,12 +652,9 @@ def run(self, mod):
         if not mod.body:
             # Nothing to do.
             return
+
         # Insert some special imports at the top of the module but after any
         # docstrings and __future__ imports.
-        aliases = [
-            ast.alias(six.moves.builtins.__name__, "@py_builtins"),
-            ast.alias("_pytest.assertion.rewrite", "@pytest_ar"),
-        ]
         doc = getattr(mod, "docstring", None)
         expect_docstring = doc is None
         if doc is not None and self.is_rewrite_disabled(doc):
@@ -684,6 +681,19 @@ def run(self, mod):
             pos += 1
         else:
             lineno = item.lineno
+        if sys.version_info >= (3, 10):
+            aliases = [
+                ast.alias(six.moves.builtins.__name__, "@py_builtins", lineno=lineno, col_offset=0),
+                ast.alias(
+                    "_pytest.assertion.rewrite", "@pytest_ar",
+                    lineno=lineno, col_offset=0
+                ),
+            ]
+        else:
+            aliases = [
+                ast.alias(six.moves.builtins.__name__, "@py_builtins"),
+                ast.alias("_pytest.assertion.rewrite", "@pytest_ar"),
+            ]
         imports = [
             ast.Import([alias], lineno=lineno, col_offset=0) for alias in aliases
         ]

From f205b1e104a446c03742726a192d92e6747f0095 Mon Sep 17 00:00:00 2001
From: Ran Benita <ran@unusedvar.com>
Date: Tue, 5 Oct 2021 09:52:09 +0300
Subject: [PATCH 2/2] rewrite: fixup end_lineno, end_col_offset of rewritten
 asserts

These are new additions in Python 3.8:
https://docs.python.org/3/whatsnew/3.8.html#ast
I'm not sure what's using them but we should set them anyway.
---
 changelog/9163.bugfix.rst        |  1 +
 src/_pytest/assertion/rewrite.py | 24 +++++++++---------------
 testing/test_assertrewrite.py    | 22 ++++++++++++++++++++++
 3 files changed, 32 insertions(+), 15 deletions(-)
 create mode 100644 changelog/9163.bugfix.rst

diff --git a/changelog/9163.bugfix.rst b/changelog/9163.bugfix.rst
new file mode 100644
index 00000000000..fb559d10fe8
--- /dev/null
+++ b/changelog/9163.bugfix.rst
@@ -0,0 +1 @@
+The end line number and end column offset are now properly set for rewritten assert statements.
diff --git a/src/_pytest/assertion/rewrite.py b/src/_pytest/assertion/rewrite.py
index 3cc17ad4526..1bd4df296fb 100644
--- a/src/_pytest/assertion/rewrite.py
+++ b/src/_pytest/assertion/rewrite.py
@@ -575,19 +575,12 @@ def _NameConstant(c):
         return ast.Name(str(c), ast.Load())
 
 
-def set_location(node, lineno, col_offset):
-    """Set node location information recursively."""
-
-    def _fix(node, lineno, col_offset):
-        if "lineno" in node._attributes:
-            node.lineno = lineno
-        if "col_offset" in node._attributes:
-            node.col_offset = col_offset
-        for child in ast.iter_child_nodes(node):
-            _fix(child, lineno, col_offset)
-
-    _fix(node, lineno, col_offset)
-    return node
+def traverse_node(node):
+    """Recursively yield node and all its children in depth-first order."""
+    yield node
+    for child in ast.iter_child_nodes(node):
+        for descendant in traverse_node(child):
+            yield descendant
 
 
 class AssertionRewriter(ast.NodeVisitor):
@@ -868,9 +861,10 @@ def visit_Assert(self, assert_):
             variables = [ast.Name(name, ast.Store()) for name in self.variables]
             clear = ast.Assign(variables, _NameConstant(None))
             self.statements.append(clear)
-        # Fix line numbers.
+        # Fix locations (line numbers/column offsets).
         for stmt in self.statements:
-            set_location(stmt, assert_.lineno, assert_.col_offset)
+            for node in traverse_node(stmt):
+                ast.copy_location(node, assert_)
         return self.statements
 
     def warn_about_none_ast(self, node, module_path, lineno):
diff --git a/testing/test_assertrewrite.py b/testing/test_assertrewrite.py
index 87dada213d6..a855c17bc30 100644
--- a/testing/test_assertrewrite.py
+++ b/testing/test_assertrewrite.py
@@ -111,6 +111,28 @@ def test_place_initial_imports(self):
             assert imp.col_offset == 0
         assert isinstance(m.body[3], ast.Expr)
 
+    def test_location_is_set(self):
+        s = textwrap.dedent(
+            """
+
+        assert False, (
+
+            "Ouch"
+          )
+
+        """
+        )
+        m = rewrite(s)
+        for node in m.body:
+            if isinstance(node, ast.Import):
+                continue
+            for n in [node, *ast.iter_child_nodes(node)]:
+                assert n.lineno == 3
+                assert n.col_offset == 0
+                if sys.version_info >= (3, 8):
+                    assert n.end_lineno == 6
+                    assert n.end_col_offset == 3
+
     def test_dont_rewrite(self):
         s = """'PYTEST_DONT_REWRITE'\nassert 14"""
         m = rewrite(s)