@@ -684,12 +684,9 @@ def run(self, mod: ast.Module) -> None:
684
684
if not mod .body :
685
685
# Nothing to do.
686
686
return
687
- # Insert some special imports at the top of the module but after any
688
- # docstrings and __future__ imports.
689
- aliases = [
690
- ast .alias ("builtins" , "@py_builtins" ),
691
- ast .alias ("_pytest.assertion.rewrite" , "@pytest_ar" ),
692
- ]
687
+
688
+ # We'll insert some special imports at the top of the module, but after any
689
+ # docstrings and __future__ imports, so first figure out where that is.
693
690
doc = getattr (mod , "docstring" , None )
694
691
expect_docstring = doc is None
695
692
if doc is not None and self .is_rewrite_disabled (doc ):
@@ -721,10 +718,27 @@ def run(self, mod: ast.Module) -> None:
721
718
lineno = item .decorator_list [0 ].lineno
722
719
else :
723
720
lineno = item .lineno
721
+ # Now actually insert the special imports.
722
+ if sys .version_info >= (3 , 10 ):
723
+ aliases = [
724
+ ast .alias ("builtins" , "@py_builtins" , lineno = lineno , col_offset = 0 ),
725
+ ast .alias (
726
+ "_pytest.assertion.rewrite" ,
727
+ "@pytest_ar" ,
728
+ lineno = lineno ,
729
+ col_offset = 0 ,
730
+ ),
731
+ ]
732
+ else :
733
+ aliases = [
734
+ ast .alias ("builtins" , "@py_builtins" ),
735
+ ast .alias ("_pytest.assertion.rewrite" , "@pytest_ar" ),
736
+ ]
724
737
imports = [
725
738
ast .Import ([alias ], lineno = lineno , col_offset = 0 ) for alias in aliases
726
739
]
727
740
mod .body [pos :pos ] = imports
741
+
728
742
# Collect asserts.
729
743
nodes : List [ast .AST ] = [mod ]
730
744
while nodes :
0 commit comments