Merge pull request #686 from Senpos/import-removal-fix

fix imports removal not including rules from previous versions
asottile · Jul 26, 2022 · 067f1e0 · 067f1e0
2 parents 35d9060 + 360e5f4
commit 067f1e0
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/pyupgrade/_plugins/imports.py b/pyupgrade/_plugins/imports.py
@@ -233,10 +233,11 @@ def _for_version(
         Mapping[tuple[str, str], str],
         Mapping[str, str],
 ]:
-    removals = {}
+    removals = collections.defaultdict(set)
     for ver, ver_removals in REMOVALS.items():
         if ver <= version:
-            removals.update(ver_removals)
+            for base, names in ver_removals.items():
+                removals[base].update(names)
 
     exact = {}
     for ver, ver_exact in REPLACE_EXACT.items():

diff --git a/tests/features/import_removals_test.py b/tests/features/import_removals_test.py
@@ -39,6 +39,7 @@ def test_import_removals_noop(s, min_version):
         ('from __future__ import division\n', (3,), ''),
         ('from __future__ import division\n', (3, 6), ''),
         ('from __future__ import (generators,)', (2, 7), ''),
+        ('from __future__ import print_function', (3, 8), ''),
         ('from builtins import map', (3,), ''),
         ('from builtins import *', (3,), ''),
         ('from six.moves import map', (3,), ''),