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] 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 000000000..fb559d10f
--- /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 3cc17ad45..1bd4df296 100644
--- a/src/_pytest/assertion/rewrite.py
+++ b/src/_pytest/assertion/rewrite.py
@@ -575,19 +575,12 @@ else:
         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 @@ class AssertionRewriter(ast.NodeVisitor):
             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 87dada213..a855c17bc 100644
--- a/testing/test_assertrewrite.py
+++ b/testing/test_assertrewrite.py
@@ -111,6 +111,28 @@ class TestAssertionRewrite(object):
             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)