python · phsilva · Feb 25, 2020
diff --git 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
@@ -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
@@ -0,0 +1,2 @@
+lib2to3: Fix duplicated atexit import when fixing multiple sys.exitfunc.
+Patch by Paulo Henrique Silva.
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		lib2to3: Fix duplicated atexit import when fixing multiple sys.exitfunc.
		Patch by Paulo Henrique Silva.