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