ยปCore Development>Code coverage>Lib/lib2to3/fixes/fix_exitfunc.py

# Python code coverage for Lib/lib2to3/fixes/fix_exitfunc.py

#countcontent
1n/a"""
2n/aConvert use of sys.exitfunc to use the atexit module.
3n/a"""
4n/a
5n/a# Author: Benjamin Peterson
6n/a
7n/afrom lib2to3 import pytree, fixer_base
8n/afrom lib2to3.fixer_util import Name, Attr, Call, Comma, Newline, syms
9n/a
10n/a
11n/aclass FixExitfunc(fixer_base.BaseFix):
12n/a keep_line_order = True
13n/a BM_compatible = True
14n/a
15n/a PATTERN = """
16n/a (
17n/a sys_import=import_name<'import'
18n/a ('sys'
19n/a |
20n/a dotted_as_names< (any ',')* 'sys' (',' any)* >
21n/a )
22n/a >
23n/a |
24n/a expr_stmt<
25n/a power< 'sys' trailer< '.' 'exitfunc' > >
26n/a '=' func=any >
27n/a )
28n/a """
29n/a
30n/a def __init__(self, *args):
31n/a super(FixExitfunc, self).__init__(*args)
32n/a
33n/a def start_tree(self, tree, filename):
34n/a super(FixExitfunc, self).start_tree(tree, filename)
35n/a self.sys_import = None
36n/a
37n/a def transform(self, node, results):
38n/a # First, find the sys import. We'll just hope it's global scope.
39n/a if "sys_import" in results:
40n/a if self.sys_import is None:
41n/a self.sys_import = results["sys_import"]
42n/a return
43n/a
44n/a func = results["func"].clone()
45n/a func.prefix = ""
46n/a register = pytree.Node(syms.power,
47n/a Attr(Name("atexit"), Name("register"))
48n/a )
49n/a call = Call(register, [func], node.prefix)
50n/a node.replace(call)
51n/a
52n/a if self.sys_import is None:
53n/a # That's interesting.
54n/a self.warning(node, "Can't find sys import; Please add an atexit "
55n/a "import at the top of your file.")
56n/a return
57n/a
58n/a # Now add an atexit import after the sys import.
59n/a names = self.sys_import.children[1]
60n/a if names.type == syms.dotted_as_names:
61n/a names.append_child(Comma())
62n/a names.append_child(Name("atexit", " "))
63n/a else:
64n/a containing_stmt = self.sys_import.parent
65n/a position = containing_stmt.children.index(self.sys_import)
66n/a stmt_container = containing_stmt.parent
67n/a new_import = pytree.Node(syms.import_name,
68n/a [Name("import"), Name("atexit", " ")]
69n/a )
70n/a new = pytree.Node(syms.simple_stmt, [new_import])
71n/a containing_stmt.insert_child(position + 1, Newline())
72n/a containing_stmt.insert_child(position + 2, new)