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