1798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson""" 2798e54004f47b8173ef70e6df0132bedceb13311Benjamin PetersonConvert use of sys.exitfunc to use the atexit module. 3798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson""" 4798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson 5798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson# Author: Benjamin Peterson 6798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson 7798e54004f47b8173ef70e6df0132bedceb13311Benjamin Petersonfrom lib2to3 import pytree, fixer_base 8798e54004f47b8173ef70e6df0132bedceb13311Benjamin Petersonfrom lib2to3.fixer_util import Name, Attr, Call, Comma, Newline, syms 9798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson 10798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson 11798e54004f47b8173ef70e6df0132bedceb13311Benjamin Petersonclass FixExitfunc(fixer_base.BaseFix): 12a81eae1fd71eed670d1c28bf940ea99b4bfa2335Benjamin Peterson keep_line_order = True 13a81eae1fd71eed670d1c28bf940ea99b4bfa2335Benjamin Peterson BM_compatible = True 14798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson 15798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson PATTERN = """ 16798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson ( 17798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson sys_import=import_name<'import' 18798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson ('sys' 19798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson | 20798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson dotted_as_names< (any ',')* 'sys' (',' any)* > 21798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson ) 22798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson > 23798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson | 24798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson expr_stmt< 25798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson power< 'sys' trailer< '.' 'exitfunc' > > 26798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson '=' func=any > 27798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson ) 28798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson """ 29798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson 30798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson def __init__(self, *args): 31798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson super(FixExitfunc, self).__init__(*args) 32798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson 33798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson def start_tree(self, tree, filename): 34798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson super(FixExitfunc, self).start_tree(tree, filename) 35798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson self.sys_import = None 36798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson 37798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson def transform(self, node, results): 3807f90476b954382e7a7ce95af696d90100a6eea1Benjamin Peterson # First, find the sys import. We'll just hope it's global scope. 39798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson if "sys_import" in results: 40798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson if self.sys_import is None: 41798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson self.sys_import = results["sys_import"] 42798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson return 43798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson 44798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson func = results["func"].clone() 45996e88a7d5dc94f060ef7be7d5235419841a492dBenjamin Peterson func.prefix = u"" 46798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson register = pytree.Node(syms.power, 47996e88a7d5dc94f060ef7be7d5235419841a492dBenjamin Peterson Attr(Name(u"atexit"), Name(u"register")) 48798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson ) 49798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson call = Call(register, [func], node.prefix) 50798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson node.replace(call) 51798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson 52798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson if self.sys_import is None: 53798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson # That's interesting. 54798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson self.warning(node, "Can't find sys import; Please add an atexit " 55798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson "import at the top of your file.") 56798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson return 57798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson 58798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson # Now add an atexit import after the sys import. 59798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson names = self.sys_import.children[1] 60798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson if names.type == syms.dotted_as_names: 61798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson names.append_child(Comma()) 62996e88a7d5dc94f060ef7be7d5235419841a492dBenjamin Peterson names.append_child(Name(u"atexit", u" ")) 63798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson else: 64798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson containing_stmt = self.sys_import.parent 65798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson position = containing_stmt.children.index(self.sys_import) 66798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson stmt_container = containing_stmt.parent 673bcc35b795862bf1ac7b8a515089843a98aa0848Benjamin Peterson new_import = pytree.Node(syms.import_name, 68996e88a7d5dc94f060ef7be7d5235419841a492dBenjamin Peterson [Name(u"import"), Name(u"atexit", u" ")] 69798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson ) 70996e88a7d5dc94f060ef7be7d5235419841a492dBenjamin Peterson new = pytree.Node(syms.simple_stmt, [new_import]) 71798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson containing_stmt.insert_child(position + 1, Newline()) 72798e54004f47b8173ef70e6df0132bedceb13311Benjamin Peterson containing_stmt.insert_child(position + 2, new) 73