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
This commit is contained in:
hauntsaninja 2021-04-11 15:38:32 -07:00
parent 9653a0e9f4
commit 8be1628042
3 changed files with 16 additions and 4 deletions

View File

@ -277,6 +277,7 @@ Sankt Petersbug
Segev Finer
Serhii Mozghovyi
Seth Junot
Shantanu Jain
Shubham Adep
Simon Gomizelj
Simon Kerr

View File

@ -0,0 +1 @@
Fixed assertion rewriting on Python 3.10.

View File

@ -684,12 +684,9 @@ class AssertionRewriter(ast.NodeVisitor):
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("builtins", "@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):
@ -721,6 +718,19 @@ class AssertionRewriter(ast.NodeVisitor):
lineno = item.decorator_list[0].lineno
else:
lineno = item.lineno
if sys.version_info >= (3, 10):
aliases = [
ast.alias("builtins", "@py_builtins", lineno=lineno, col_offset=0),
ast.alias(
"_pytest.assertion.rewrite", "@pytest_ar",
lineno=lineno, col_offset=0
),
]
else:
aliases = [
ast.alias("builtins", "@py_builtins"),
ast.alias("_pytest.assertion.rewrite", "@pytest_ar"),
]
imports = [
ast.Import([alias], lineno=lineno, col_offset=0) for alias in aliases
]