| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172 |
- """
- Convert use of sys.exitfunc to use the atexit module.
- """
- # Author: Benjamin Peterson
- from lib2to3 import pytree, fixer_base
- from lib2to3.fixer_util import Name, Attr, Call, Comma, Newline, syms
- class FixExitfunc(fixer_base.BaseFix):
- keep_line_order = True
- BM_compatible = True
- PATTERN = """
- (
- sys_import=import_name<'import'
- ('sys'
- |
- dotted_as_names< (any ',')* 'sys' (',' any)* >
- )
- >
- |
- expr_stmt<
- power< 'sys' trailer< '.' 'exitfunc' > >
- '=' func=any >
- )
- """
- def __init__(self, *args):
- super(FixExitfunc, self).__init__(*args)
- def start_tree(self, tree, filename):
- super(FixExitfunc, self).start_tree(tree, filename)
- self.sys_import = None
- def transform(self, node, results):
- # First, find the sys import. We'll just hope it's global scope.
- if "sys_import" in results:
- if self.sys_import is None:
- self.sys_import = results["sys_import"]
- return
- func = results["func"].clone()
- func.prefix = ""
- register = pytree.Node(syms.power,
- Attr(Name("atexit"), Name("register"))
- )
- call = Call(register, [func], node.prefix)
- node.replace(call)
- if self.sys_import is None:
- # That's interesting.
- self.warning(node, "Can't find sys import; Please add an atexit "
- "import at the top of your file.")
- return
- # Now add an atexit import after the sys import.
- names = self.sys_import.children[1]
- if names.type == syms.dotted_as_names:
- names.append_child(Comma())
- names.append_child(Name("atexit", " "))
- else:
- containing_stmt = self.sys_import.parent
- position = containing_stmt.children.index(self.sys_import)
- stmt_container = containing_stmt.parent
- new_import = pytree.Node(syms.import_name,
- [Name("import"), Name("atexit", " ")]
- )
- new = pytree.Node(syms.simple_stmt, [new_import])
- containing_stmt.insert_child(position + 1, Newline())
- containing_stmt.insert_child(position + 2, new)
|
