2 Convert use of sys.exitfunc to use the atexit module.
5 # Author: Benjamin Peterson
7 from lib2to3
import pytree
, fixer_base
8 from lib2to3
.fixer_util
import Name
, Attr
, Call
, Comma
, Newline
, syms
11 class FixExitfunc(fixer_base
.BaseFix
):
15 sys_import=import_name<'import'
18 dotted_as_names< (any ',')* 'sys' (',' any)* >
23 power< 'sys' trailer< '.' 'exitfunc' > >
28 def __init__(self
, *args
):
29 super(FixExitfunc
, self
).__init
__(*args
)
31 def start_tree(self
, tree
, filename
):
32 super(FixExitfunc
, self
).start_tree(tree
, filename
)
33 self
.sys_import
= None
35 def transform(self
, node
, results
):
36 # First, find a the sys import. We'll just hope it's global scope.
37 if "sys_import" in results
:
38 if self
.sys_import
is None:
39 self
.sys_import
= results
["sys_import"]
42 func
= results
["func"].clone()
44 register
= pytree
.Node(syms
.power
,
45 Attr(Name(u
"atexit"), Name(u
"register"))
47 call
= Call(register
, [func
], node
.prefix
)
50 if self
.sys_import
is None:
52 self
.warning(node
, "Can't find sys import; Please add an atexit "
53 "import at the top of your file.")
56 # Now add an atexit import after the sys import.
57 names
= self
.sys_import
.children
[1]
58 if names
.type == syms
.dotted_as_names
:
59 names
.append_child(Comma())
60 names
.append_child(Name(u
"atexit", u
" "))
62 containing_stmt
= self
.sys_import
.parent
63 position
= containing_stmt
.children
.index(self
.sys_import
)
64 stmt_container
= containing_stmt
.parent
65 new_import
= pytree
.Node(syms
.import_name
,
66 [Name(u
"import"), Name(u
"atexit", u
" ")]
68 new
= pytree
.Node(syms
.simple_stmt
, [new_import
])
69 containing_stmt
.insert_child(position
+ 1, Newline())
70 containing_stmt
.insert_child(position
+ 2, new
)