Remap ChainMap, Counter, and OrderedDict imports to collections (#3392)

astral-sh · Mar 7, 2023 · 3f04def · 3f04def
1 parent 9817775
commit 3f04def
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 9 deletions.
diff --git a/crates/ruff/resources/test/fixtures/pyupgrade/UP035.py b/crates/ruff/resources/test/fixtures/pyupgrade/UP035.py
@@ -41,7 +41,7 @@
         Good,
     )
 
-from typing import Callable, Match, Pattern, List
+from typing import Callable, Match, Pattern, List, OrderedDict
 
 if True: from collections import (
     Mapping, Counter)

diff --git a/crates/ruff/src/rules/pyupgrade/rules/import_replacements.rs b/crates/ruff/src/rules/pyupgrade/rules/import_replacements.rs
@@ -134,11 +134,9 @@ const TYPING_TO_COLLECTIONS_ABC_39: &[&str] = &[
     "AsyncIterator",
     "Awaitable",
     "ByteString",
-    "ChainMap",
     "Collection",
     "Container",
     "Coroutine",
-    "Counter",
     "Generator",
     "Hashable",
     "ItemsView",
@@ -156,6 +154,9 @@ const TYPING_TO_COLLECTIONS_ABC_39: &[&str] = &[
     "ValuesView",
 ];
 
+// Members of `typing` that have been moved to `collections`.
+const TYPING_TO_COLLECTIONS_39: &[&str] = &["ChainMap", "Counter", "OrderedDict"];
+
 // Members of `typing` that have been moved to `typing.re`.
 const TYPING_TO_RE_39: &[&str] = &["Match", "Pattern"];
 
@@ -301,6 +302,15 @@ impl<'a> ImportReplacer<'a> {
                     replacements.push(replacement);
                 }
 
+                // `typing` to `collections`
+                let mut typing_to_collections = vec![];
+                if self.version >= PythonVersion::Py39 {
+                    typing_to_collections.extend(TYPING_TO_COLLECTIONS_39);
+                }
+                if let Some(replacement) = self.try_replace(&typing_to_collections, "collections") {
+                    replacements.push(replacement);
+                }
+
                 // `typing` to `re`
                 let mut typing_to_re = vec![];
                 if self.version >= PythonVersion::Py39 {

diff --git a/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP035.py.snap b/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP035.py.snap
@@ -310,15 +310,36 @@ expression: diagnostics
     column: 0
   end_location:
     row: 44
-    column: 49
+    column: 62
   fix:
-    content: "from typing import Match, Pattern, List\nfrom collections.abc import Callable"
+    content: "from typing import Match, Pattern, List, OrderedDict\nfrom collections.abc import Callable"
     location:
       row: 44
       column: 0
     end_location:
       row: 44
-      column: 49
+      column: 62
+  parent: ~
+- kind:
+    ImportReplacements:
+      module: collections
+      members:
+        - OrderedDict
+      fixable: true
+  location:
+    row: 44
+    column: 0
+  end_location:
+    row: 44
+    column: 62
+  fix:
+    content: "from typing import Callable, Match, Pattern, List\nfrom collections import OrderedDict"
+    location:
+      row: 44
+      column: 0
+    end_location:
+      row: 44
+      column: 62
   parent: ~
 - kind:
     ImportReplacements:
@@ -332,15 +353,15 @@ expression: diagnostics
     column: 0
   end_location:
     row: 44
-    column: 49
+    column: 62
   fix:
-    content: "from typing import Callable, List\nfrom re import Match, Pattern"
+    content: "from typing import Callable, List, OrderedDict\nfrom re import Match, Pattern"
     location:
       row: 44
       column: 0
     end_location:
       row: 44
-      column: 49
+      column: 62
   parent: ~
 - kind:
     ImportReplacements: