Pyupgrade: import mock to from unittest import mock (#1488)

README.md

@@ -688,6 +688,7 @@ For more, see [pyupgrade](https://pypi.org/project/pyupgrade/3.2.0/) on PyPI.
 | UP023 | RewriteCElementTree | `cElementTree` is deprecated, use `ElementTree` | 🛠 |
 | UP024 | OSErrorAlias | Replace aliased errors with `OSError` | 🛠 |
 | UP025 | RewriteUnicodeLiteral | Remove unicode literals from strings | 🛠 |
+| UP026 | RewriteMockImport | `mock` is deprecated, use `unittest.mock` | 🛠 |
 
 ### pep8-naming (N)
 

foo.py

@@ -0,0 +1,13 @@
+if True:
+    from unittest import mock as foo
+    from unittest import mock as bar
+    from unittest import mock
+    import os
+    from unittest import mock as foo
+    from unittest import mock as bar
+    from unittest import mock
+
+if True:
+    from unittest import mock as foo
+    from unittest import mock as bar
+    from unittest import mock

resources/test/fixtures/pyupgrade/UP026.py

@@ -0,0 +1,67 @@
+# 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
+
+# Mock should go on a new line as `from unittest import mock`
+import mock, mock
+
+# Mock should go on a new line as `from unittest import mock as foo`
+import mock as foo
+
+# Mock should go on a new line as `from unittest import mock as foo`
+from mock import mock as foo
+
+if True:
+    # This should yield multiple, aliased imports.
+    import mock as foo, mock as bar, mock
+
+    # This should yield multiple, aliased imports, and preserve `os`.
+    import mock as foo, mock as bar, mock, os
+
+if True:
+    # This should yield multiple, aliased imports.
+    from mock import mock as foo, mock as bar, mock

ruff.schema.json

@@ -911,6 +911,7 @@
         "UP023",
         "UP024",
         "UP025",
+        "UP026",
         "W",
         "W2",
         "W29",

src/checkers/ast.rs

@@ -643,6 +643,9 @@ where
                 if self.settings.enabled.contains(&CheckCode::UP023) {
                     pyupgrade::plugins::replace_c_element_tree(self, stmt);
                 }
+                if self.settings.enabled.contains(&CheckCode::UP026) {
+                    pyupgrade::plugins::rewrite_mock_import(self, stmt);
+                }
 
                 for alias in names {
                     if alias.node.name.contains('.') && alias.node.asname.is_none() {
@@ -854,6 +857,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

@@ -238,6 +238,7 @@ pub enum CheckCode {
     UP023,
     UP024,
     UP025,
+    UP026,
     // pydocstyle
     D100,
     D101,
@@ -910,6 +911,7 @@ pub enum CheckKind {
     RewriteCElementTree,
     OSErrorAlias(Option<String>),
     RewriteUnicodeLiteral,
+    RewriteMockImport,
     // pydocstyle
     BlankLineAfterLastSection(String),
     BlankLineAfterSection(String),
@@ -1305,6 +1307,7 @@ impl CheckCode {
             CheckCode::UP023 => CheckKind::RewriteCElementTree,
             CheckCode::UP024 => CheckKind::OSErrorAlias(None),
             CheckCode::UP025 => CheckKind::RewriteUnicodeLiteral,
+            CheckCode::UP026 => CheckKind::RewriteMockImport,
             // pydocstyle
             CheckCode::D100 => CheckKind::PublicModule,
             CheckCode::D101 => CheckKind::PublicClass,
@@ -1738,6 +1741,7 @@ impl CheckCode {
             CheckCode::UP023 => CheckCategory::Pyupgrade,
             CheckCode::UP024 => CheckCategory::Pyupgrade,
             CheckCode::UP025 => CheckCategory::Pyupgrade,
+            CheckCode::UP026 => CheckCategory::Pyupgrade,
             CheckCode::W292 => CheckCategory::Pycodestyle,
             CheckCode::W605 => CheckCategory::Pycodestyle,
             CheckCode::YTT101 => CheckCategory::Flake82020,
@@ -1961,6 +1965,7 @@ impl CheckKind {
             CheckKind::RewriteCElementTree => &CheckCode::UP023,
             CheckKind::OSErrorAlias(..) => &CheckCode::UP024,
             CheckKind::RewriteUnicodeLiteral => &CheckCode::UP025,
+            CheckKind::RewriteMockImport => &CheckCode::UP026,
             // pydocstyle
             CheckKind::BlankLineAfterLastSection(..) => &CheckCode::D413,
             CheckKind::BlankLineAfterSection(..) => &CheckCode::D410,
@@ -2722,6 +2727,7 @@ impl CheckKind {
             }
             CheckKind::OSErrorAlias(..) => "Replace aliased errors with `OSError`".to_string(),
             CheckKind::RewriteUnicodeLiteral => "Remove unicode literals from strings".to_string(),
+            CheckKind::RewriteMockImport => "`mock` is deprecated, use `unittest.mock`".to_string(),
             // pydocstyle
             CheckKind::FitsOnOneLine => "One-line docstring should fit on one line".to_string(),
             CheckKind::BlankLineAfterSummary => {
@@ -3190,6 +3196,7 @@ impl CheckKind {
                 | CheckKind::ReplaceStdoutStderr
                 | CheckKind::ReplaceUniversalNewlines
                 | CheckKind::RewriteCElementTree
+                | CheckKind::RewriteMockImport
                 | CheckKind::RewriteUnicodeLiteral
                 | CheckKind::SectionNameEndsInColon(..)
                 | CheckKind::SectionNotOverIndented(..)
@@ -3302,6 +3309,7 @@ impl CheckKind {
             }
             CheckKind::RewriteCElementTree => Some("Replace with `ElementTree`".to_string()),
             CheckKind::RewriteUnicodeLiteral => Some("Remove unicode prefix".to_string()),
+            CheckKind::RewriteMockImport => Some("Import from `unittest.mock` instead".to_string()),
             CheckKind::NewLineAfterSectionName(name) => {
                 Some(format!("Add newline after \"{name}\""))
             }

src/checks_gen.rs

@@ -542,6 +542,7 @@ pub enum CheckCodePrefix {
     UP023,
     UP024,
     UP025,
+    UP026,
     W,
     W2,
     W29,
@@ -776,6 +777,7 @@ impl CheckCodePrefix {
                 CheckCode::UP023,
                 CheckCode::UP024,
                 CheckCode::UP025,
+                CheckCode::UP026,
                 CheckCode::D100,
                 CheckCode::D101,
                 CheckCode::D102,
@@ -2456,6 +2458,7 @@ impl CheckCodePrefix {
                     CheckCode::UP023,
                     CheckCode::UP024,
                     CheckCode::UP025,
+                    CheckCode::UP026,
                 ]
             }
             CheckCodePrefix::U0 => {
@@ -2490,6 +2493,7 @@ impl CheckCodePrefix {
                     CheckCode::UP023,
                     CheckCode::UP024,
                     CheckCode::UP025,
+                    CheckCode::UP026,
                 ]
             }
             CheckCodePrefix::U00 => {
@@ -2708,6 +2712,7 @@ impl CheckCodePrefix {
                 CheckCode::UP023,
                 CheckCode::UP024,
                 CheckCode::UP025,
+                CheckCode::UP026,
             ],
             CheckCodePrefix::UP0 => vec![
                 CheckCode::UP001,
@@ -2734,6 +2739,7 @@ impl CheckCodePrefix {
                 CheckCode::UP023,
                 CheckCode::UP024,
                 CheckCode::UP025,
+                CheckCode::UP026,
             ],
             CheckCodePrefix::UP00 => vec![
                 CheckCode::UP001,
@@ -2782,13 +2788,15 @@ impl CheckCodePrefix {
                 CheckCode::UP023,
                 CheckCode::UP024,
                 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::UP024 => vec![CheckCode::UP024],
             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::UP023 => SuffixLength::Three,
             CheckCodePrefix::UP024 => SuffixLength::Three,
             CheckCodePrefix::UP025 => SuffixLength::Three,
+            CheckCodePrefix::UP026 => SuffixLength::Three,
             CheckCodePrefix::W => SuffixLength::Zero,
             CheckCodePrefix::W2 => SuffixLength::One,
             CheckCodePrefix::W29 => SuffixLength::Two,

src/cst/matchers.rs

@@ -1,5 +1,5 @@
 use anyhow::{bail, Result};
-use libcst_native::{Expr, Module, SmallStatement, Statement};
+use libcst_native::{Expr, Import, ImportFrom, Module, SmallStatement, Statement};
 
 pub fn match_module(module_text: &str) -> Result<Module> {
     match libcst_native::parse_module(module_text, None) {
@@ -19,3 +19,27 @@ pub fn match_expr<'a, 'b>(module: &'a mut Module<'b>) -> Result<&'a mut Expr<'b>
         bail!("Expected Statement::Simple")
     }
 }
+
+pub fn match_import<'a, 'b>(module: &'a mut Module<'b>) -> Result<&'a mut Import<'b>> {
+    if let Some(Statement::Simple(expr)) = module.body.first_mut() {
+        if let Some(SmallStatement::Import(expr)) = expr.body.first_mut() {
+            Ok(expr)
+        } else {
+            bail!("Expected SmallStatement::Expr")
+        }
+    } else {
+        bail!("Expected Statement::Simple")
+    }
+}
+
+pub fn match_import_from<'a, 'b>(module: &'a mut Module<'b>) -> Result<&'a mut ImportFrom<'b>> {
+    if let Some(Statement::Simple(expr)) = module.body.first_mut() {
+        if let Some(SmallStatement::ImportFrom(expr)) = expr.body.first_mut() {
+            Ok(expr)
+        } else {
+            bail!("Expected SmallStatement::Expr")
+        }
+    } else {
+        bail!("Expected Statement::Simple")
+    }
+}

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

@@ -46,6 +46,7 @@ mod tests {
     #[test_case(CheckCode::UP024, Path::new("UP024_1.py"); "UP024_1")]
     #[test_case(CheckCode::UP024, Path::new("UP024_2.py"); "UP024_2")]
     #[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 checks = test_path(

src/pyupgrade/plugins/mod.rs

@@ -10,6 +10,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;
@@ -34,6 +35,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;