Pyupgrade: import mock to from unittest import mock

README.md

@@ -685,6 +685,7 @@ For more, see [pyupgrade](https://pypi.org/project/pyupgrade/3.2.0/) on PyPI.
 | UP022 | ReplaceStdoutStderr | Sending stdout and stderr to pipe is deprecated, use `capture_output` | 🛠 |
 | UP023 | RewriteCElementTree | `cElementTree` is deprecated, use `ElementTree` | 🛠 |
 | UP025 | RewriteUnicodeLiteral | Remove unicode literals from strings | 🛠 |
+| UP026 | RewriteMockImport | `mock` is deprecated, use `unittest.mock` | 🛠 |
 
 ### pep8-naming (N)
 
@@ -2805,6 +2806,28 @@ convention = "google"
 
 ### `pyupgrade`
 
+#### [`keep-mock`](#keep-mock)
+
+Whether to avoid PEP 585 (`List[int]` -> `list[int]`) and PEP 604
+(`Optional[str]` -> `str | None`) rewrites even if a file imports `from
+__future__ import annotations`. Note that this setting is only
+applicable when the target Python version is below 3.9 and 3.10
+respectively.
+
+**Default value**: `false`
+
+**Type**: `bool`
+
+**Example usage**:
+
+```toml
+[tool.ruff.pyupgrade]
+# Does not replace `mock` imports with `from unittest import mock`.
+keep-mock = true
+```
+
+---
+
 #### [`keep-runtime-typing`](#keep-runtime-typing)
 
 Whether to avoid PEP 585 (`List[int]` -> `list[int]`) and PEP 604

resources/test/fixtures/pyupgrade/UP026.py

@@ -0,0 +1,48 @@
+# These should be changed
+if True:
+    import mock
+
+if True:
+    import mock, sys
+
+# This goes to from unitest import mock
+import mock.mock
+
+# Mock should go on a new line as `from unittest import mock`
+import contextlib, mock, sys
+
+# Mock should go on a new line as `from unittest import mock`
+import mock, sys
+x = "This code should be preserved one line below the mock"
+
+# Mock should go on a new line as `from unittest import mock`
+from mock import mock
+
+# Should keep trailing comma
+from mock import (
+    mock,
+    a,
+    b,
+    c,
+)
+
+# Should not get a trailing comma
+from mock import (
+    mock,
+    a,
+    b,
+    c
+)
+
+if True:
+    if False:
+        from mock import (
+            mock,
+            a,
+            b,
+            c
+        )
+
+
+# These should not change:
+import os, io

ruff.schema.json

@@ -912,6 +912,7 @@
         "UP022",
         "UP023",
         "UP025",
+        "UP026",
         "W",
         "W2",
         "W29",
@@ -1261,6 +1262,13 @@
     "PyUpgradeOptions": {
       "type": "object",
       "properties": {
+        "keep-mock": {
+          "description": "Whether to avoid PEP 585 (`List[int]` -> `list[int]`) and PEP 604 (`Optional[str]` -> `str | None`) rewrites even if a file imports `from __future__ import annotations`. Note that this setting is only applicable when the target Python version is below 3.9 and 3.10 respectively.",
+          "type": [
+            "boolean",
+            "null"
+          ]
+        },
         "keep-runtime-typing": {
           "description": "Whether to avoid PEP 585 (`List[int]` -> `list[int]`) and PEP 604 (`Optional[str]` -> `str | None`) rewrites even if a file imports `from __future__ import annotations`. Note that this setting is only applicable when the target Python version is below 3.9 and 3.10 respectively.",
           "type": [

src/checkers/ast.rs

@@ -652,6 +652,11 @@ where
                 if self.settings.enabled.contains(&CheckCode::UP023) {
                     pyupgrade::plugins::replace_c_element_tree(self, stmt);
                 }
+                if self.settings.enabled.contains(&CheckCode::UP026)
+                    && !self.settings.pyupgrade.keep_mock
+                {
+                    pyupgrade::plugins::rewrite_mock_import(self, stmt);
+                }
 
                 for alias in names {
                     if alias.node.name.contains('.') && alias.node.asname.is_none() {
@@ -863,6 +868,9 @@ where
                         pyupgrade::plugins::unnecessary_future_import(self, stmt, names);
                     }
                 }
+                if self.settings.enabled.contains(&CheckCode::UP026) {
+                    pyupgrade::plugins::rewrite_mock_import(self, stmt);
+                }
 
                 if self.settings.enabled.contains(&CheckCode::TID251) {
                     if let Some(module) = module {

src/checks.rs

@@ -237,6 +237,7 @@ pub enum CheckCode {
     UP022,
     UP023,
     UP025,
+    UP026,
     // pydocstyle
     D100,
     D101,
@@ -868,6 +869,7 @@ pub enum CheckKind {
     ReplaceStdoutStderr,
     RewriteCElementTree,
     RewriteUnicodeLiteral,
+    RewriteMockImport,
     // pydocstyle
     BlankLineAfterLastSection(String),
     BlankLineAfterSection(String),
@@ -1263,6 +1265,7 @@ impl CheckCode {
             CheckCode::UP022 => CheckKind::ReplaceStdoutStderr,
             CheckCode::UP023 => CheckKind::RewriteCElementTree,
             CheckCode::UP025 => CheckKind::RewriteUnicodeLiteral,
+            CheckCode::UP026 => CheckKind::RewriteMockImport,
             // pydocstyle
             CheckCode::D100 => CheckKind::PublicModule,
             CheckCode::D101 => CheckKind::PublicClass,
@@ -1696,6 +1699,7 @@ impl CheckCode {
             CheckCode::UP022 => CheckCategory::Pyupgrade,
             CheckCode::UP023 => CheckCategory::Pyupgrade,
             CheckCode::UP025 => CheckCategory::Pyupgrade,
+            CheckCode::UP026 => CheckCategory::Pyupgrade,
             CheckCode::W292 => CheckCategory::Pycodestyle,
             CheckCode::W605 => CheckCategory::Pycodestyle,
             CheckCode::YTT101 => CheckCategory::Flake82020,
@@ -1919,6 +1923,7 @@ impl CheckKind {
             CheckKind::ReplaceStdoutStderr => &CheckCode::UP022,
             CheckKind::RewriteCElementTree => &CheckCode::UP023,
             CheckKind::RewriteUnicodeLiteral => &CheckCode::UP025,
+            CheckKind::RewriteMockImport => &CheckCode::UP026,
             // pydocstyle
             CheckKind::BlankLineAfterLastSection(..) => &CheckCode::D413,
             CheckKind::BlankLineAfterSection(..) => &CheckCode::D410,
@@ -2671,6 +2676,7 @@ impl CheckKind {
                 "`cElementTree` is deprecated, use `ElementTree`".to_string()
             }
             CheckKind::RewriteUnicodeLiteral => "Remove unicode literals from strings".to_string(),
+            CheckKind::RewriteMockImport => "`mock` is deprecated, use `unittest.mock`".to_string(),
             CheckKind::ConvertNamedTupleFunctionalToClass(name) => {
                 format!("Convert `{name}` from `NamedTuple` functional to class syntax")
             }
@@ -3120,6 +3126,7 @@ impl CheckKind {
                 | CheckKind::ReplaceStdoutStderr
                 | CheckKind::RewriteCElementTree
                 | CheckKind::RewriteUnicodeLiteral
+                | CheckKind::RewriteMockImport
                 | CheckKind::NewLineAfterSectionName(..)
                 | CheckKind::NoBlankLineAfterFunction(..)
                 | CheckKind::NoBlankLineBeforeClass(..)

src/checks_gen.rs

@@ -543,6 +543,7 @@ pub enum CheckCodePrefix {
     UP022,
     UP023,
     UP025,
+    UP026,
     W,
     W2,
     W29,
@@ -777,6 +778,7 @@ impl CheckCodePrefix {
                 CheckCode::UP022,
                 CheckCode::UP023,
                 CheckCode::UP025,
+                CheckCode::UP026,
                 CheckCode::D100,
                 CheckCode::D101,
                 CheckCode::D102,
@@ -2460,6 +2462,7 @@ impl CheckCodePrefix {
                     CheckCode::UP022,
                     CheckCode::UP023,
                     CheckCode::UP025,
+                    CheckCode::UP026,
                 ]
             }
             CheckCodePrefix::U0 => {
@@ -2493,6 +2496,7 @@ impl CheckCodePrefix {
                     CheckCode::UP022,
                     CheckCode::UP023,
                     CheckCode::UP025,
+                    CheckCode::UP026,
                 ]
             }
             CheckCodePrefix::U00 => {
@@ -2710,6 +2714,7 @@ impl CheckCodePrefix {
                 CheckCode::UP022,
                 CheckCode::UP023,
                 CheckCode::UP025,
+                CheckCode::UP026,
             ],
             CheckCodePrefix::UP0 => vec![
                 CheckCode::UP001,
@@ -2735,6 +2740,7 @@ impl CheckCodePrefix {
                 CheckCode::UP022,
                 CheckCode::UP023,
                 CheckCode::UP025,
+                CheckCode::UP026,
             ],
             CheckCodePrefix::UP00 => vec![
                 CheckCode::UP001,
@@ -2782,12 +2788,14 @@ impl CheckCodePrefix {
                 CheckCode::UP022,
                 CheckCode::UP023,
                 CheckCode::UP025,
+                CheckCode::UP026,
             ],
             CheckCodePrefix::UP020 => vec![CheckCode::UP020],
             CheckCodePrefix::UP021 => vec![CheckCode::UP021],
             CheckCodePrefix::UP022 => vec![CheckCode::UP022],
             CheckCodePrefix::UP023 => vec![CheckCode::UP023],
             CheckCodePrefix::UP025 => vec![CheckCode::UP025],
+            CheckCodePrefix::UP026 => vec![CheckCode::UP026],
             CheckCodePrefix::W => vec![CheckCode::W292, CheckCode::W605],
             CheckCodePrefix::W2 => vec![CheckCode::W292],
             CheckCodePrefix::W29 => vec![CheckCode::W292],
@@ -3362,6 +3370,7 @@ impl CheckCodePrefix {
             CheckCodePrefix::UP022 => SuffixLength::Three,
             CheckCodePrefix::UP023 => SuffixLength::Three,
             CheckCodePrefix::UP025 => SuffixLength::Three,
+            CheckCodePrefix::UP026 => SuffixLength::Three,
             CheckCodePrefix::W => SuffixLength::Zero,
             CheckCodePrefix::W2 => SuffixLength::One,
             CheckCodePrefix::W29 => SuffixLength::Two,

src/isort/mod.rs

@@ -23,12 +23,12 @@ use crate::SourceCodeLocator;
 mod categorize;
 mod comments;
 pub mod format;
-mod helpers;
+pub mod helpers;
 pub mod plugins;
 pub mod settings;
 mod sorting;
 pub mod track;
-mod types;
+pub mod types;
 
 #[derive(Debug)]
 pub struct AnnotatedAliasData<'a> {

src/pyupgrade/mod.rs

@@ -43,6 +43,7 @@ mod tests {
     #[test_case(CheckCode::UP022, Path::new("UP022.py"); "UP022")]
     #[test_case(CheckCode::UP023, Path::new("UP023.py"); "UP023")]
     #[test_case(CheckCode::UP025, Path::new("UP025.py"); "UP025")]
+    #[test_case(CheckCode::UP026, Path::new("UP026.py"); "UP026")]
     fn checks(check_code: CheckCode, path: &Path) -> Result<()> {
         let snapshot = format!("{}_{}", check_code.as_ref(), path.to_string_lossy());
         let mut checks = test_path(

src/pyupgrade/plugins/mod.rs

@@ -9,6 +9,7 @@ pub use remove_six_compat::remove_six_compat;
 pub use replace_stdout_stderr::replace_stdout_stderr;
 pub use replace_universal_newlines::replace_universal_newlines;
 pub use rewrite_c_element_tree::replace_c_element_tree;
+pub use rewrite_mock_import::rewrite_mock_import;
 pub use rewrite_unicode_literal::rewrite_unicode_literal;
 pub use super_call_with_parameters::super_call_with_parameters;
 pub use type_of_primitive::type_of_primitive;
@@ -32,6 +33,7 @@ mod remove_six_compat;
 mod replace_stdout_stderr;
 mod replace_universal_newlines;
 mod rewrite_c_element_tree;
+mod rewrite_mock_import;
 mod rewrite_unicode_literal;
 mod super_call_with_parameters;
 mod type_of_primitive;