diff --git a/Lib/lib2to3/fixes/fix_exitfunc.py b/Lib/lib2to3/fixes/fix_exitfunc.py index 2e47887afead368..452d99f787c154f 100644 --- a/Lib/lib2to3/fixes/fix_exitfunc.py +++ b/Lib/lib2to3/fixes/fix_exitfunc.py @@ -5,7 +5,7 @@ # Author: Benjamin Peterson from lib2to3 import pytree, fixer_base -from lib2to3.fixer_util import Name, Attr, Call, Comma, Newline, syms +from lib2to3.fixer_util import Name, Attr, Call, Comma, Newline, syms, does_tree_import class FixExitfunc(fixer_base.BaseFix): @@ -55,6 +55,10 @@ def transform(self, node, results): "import at the top of your file.") return + # Do not add import if already present (multiple sys.atexit's) + if does_tree_import(None, "atexit", self.sys_import.parent): + return + # Now add an atexit import after the sys import. names = self.sys_import.children[1] if names.type == syms.dotted_as_names: @@ -63,7 +67,6 @@ def transform(self, node, results): 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", " ")] ) diff --git a/Lib/lib2to3/tests/test_fixers.py b/Lib/lib2to3/tests/test_fixers.py index 121ebe68e5402ba..f1cb4258287bdfc 100644 --- a/Lib/lib2to3/tests/test_fixers.py +++ b/Lib/lib2to3/tests/test_fixers.py @@ -4527,6 +4527,24 @@ def test_simple(self): """ self.check(b, a) + def test_multiple(self): + b = """ + import sys + if 42: + sys.exitfunc = my_atexit + else: + sys.exitfunc = my_otheratexit + """ + a = """ + import sys + import atexit + if 42: + atexit.register(my_atexit) + else: + atexit.register(my_otheratexit) + """ + self.check(b, a) + def test_names_import(self): b = """ import sys, crumbs diff --git a/Misc/NEWS.d/next/Tools-Demos/2020-02-25-01-08-58.bpo-39683.-fbPik.rst b/Misc/NEWS.d/next/Tools-Demos/2020-02-25-01-08-58.bpo-39683.-fbPik.rst new file mode 100644 index 000000000000000..e0320f3cbf3aead --- /dev/null +++ b/Misc/NEWS.d/next/Tools-Demos/2020-02-25-01-08-58.bpo-39683.-fbPik.rst @@ -0,0 +1,2 @@ +lib2to3: Fix duplicated atexit import when fixing multiple sys.exitfunc. +Patch by Paulo Henrique Silva.