Lib/lib2to3/fixes/fix_exitfunc.py

   1 """
   2 Convert use of sys.exitfunc to use the atexit module.
   3 """
   4
   5 # Author: Benjamin Peterson
   6
   7 from lib2to3 import pytree, fixer_base
   8 from lib2to3.fixer_util import Name, Attr, Call, Comma, Newline, syms
   9
  10
  11 class FixExitfunc(fixer_base.BaseFix):
  12
  13     PATTERN = """
  14               (
  15                   sys_import=import_name<'import'
  16                       ('sys'
  17                       |
  18                       dotted_as_names< (any ',')* 'sys' (',' any)* >
  19                       )
  20                   >
  21               |
  22                   expr_stmt<
  23                       power< 'sys' trailer< '.' 'exitfunc' > >
  24                   '=' func=any >
  25               )
  26               """
  27
  28     def __init__(self, *args):
  29         super(FixExitfunc, self).__init__(*args)
  30
  31     def start_tree(self, tree, filename):
  32         super(FixExitfunc, self).start_tree(tree, filename)
  33         self.sys_import = None
  34
  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"]
  40             return
  41
  42         func = results["func"].clone()
  43         func.prefix = ""
  44         register = pytree.Node(syms.power,
  45                                Attr(Name("atexit"), Name("register"))
  46                                )
  47         call = Call(register, [func], node.prefix)
  48         node.replace(call)
  49
  50         if self.sys_import is None:
  51             # That's interesting.
  52             self.warning(node, "Can't find sys import; Please add an atexit "
  53                              "import at the top of your file.")
  54             return
  55
  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("atexit", " "))
  61         else:
  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("import"), Name("atexit", " ")]
  67                               )
  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)