Rewrite mock.mock attribute accesses (#1533)

astral-sh · Jan 1, 2023 · f1a183c · f1a183c
1 parent 509c6d5
commit f1a183c
Show file tree

Hide file tree

Showing 7 changed files with 120 additions and 32 deletions.
diff --git a/foo.py b/foo.py
@@ -0,0 +1,3 @@
+import mock.mock
+
+x = mock.mock.Mock()
diff --git a/resources/test/fixtures/pyupgrade/UP026.py b/resources/test/fixtures/pyupgrade/UP026.py
@@ -65,3 +65,10 @@
 if True:
     # This should yield multiple, aliased imports.
     from mock import mock as foo, mock as bar, mock
+
+
+# This should be unchanged.
+x = mock.Mock()
+
+# This should change to `mock.Mock`().
+x = mock.mock.Mock()
diff --git a/src/checkers/ast.rs b/src/checkers/ast.rs
@@ -1606,6 +1606,10 @@ where
                 if self.settings.enabled.contains(&CheckCode::UP019) {
                     pyupgrade::plugins::typing_text_str_alias(self, expr);
                 }
+                if self.settings.enabled.contains(&CheckCode::UP026) {
+                    pyupgrade::plugins::rewrite_mock_attribute(self, expr);
+                }
+
                 if self.settings.enabled.contains(&CheckCode::YTT202) {
                     flake8_2020::plugins::name_or_attribute(self, expr);
                 }

diff --git a/src/checks.rs b/src/checks.rs
@@ -705,6 +705,12 @@ pub struct UnusedCodes {
     pub unmatched: Vec<String>,
 }
 
+#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub enum MockReference {
+    Import,
+    Attribute,
+}
+
 #[derive(AsRefStr, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub enum CheckKind {
     // pycodestyle errors
@@ -911,7 +917,7 @@ pub enum CheckKind {
     RewriteCElementTree,
     OSErrorAlias(Option<String>),
     RewriteUnicodeLiteral,
-    RewriteMockImport,
+    RewriteMockImport(MockReference),
     // pydocstyle
     BlankLineAfterLastSection(String),
     BlankLineAfterSection(String),
@@ -1307,7 +1313,7 @@ impl CheckCode {
             CheckCode::UP023 => CheckKind::RewriteCElementTree,
             CheckCode::UP024 => CheckKind::OSErrorAlias(None),
             CheckCode::UP025 => CheckKind::RewriteUnicodeLiteral,
-            CheckCode::UP026 => CheckKind::RewriteMockImport,
+            CheckCode::UP026 => CheckKind::RewriteMockImport(MockReference::Import),
             // pydocstyle
             CheckCode::D100 => CheckKind::PublicModule,
             CheckCode::D101 => CheckKind::PublicClass,
@@ -1965,7 +1971,7 @@ impl CheckKind {
             CheckKind::RewriteCElementTree => &CheckCode::UP023,
             CheckKind::OSErrorAlias(..) => &CheckCode::UP024,
             CheckKind::RewriteUnicodeLiteral => &CheckCode::UP025,
-            CheckKind::RewriteMockImport => &CheckCode::UP026,
+            CheckKind::RewriteMockImport(..) => &CheckCode::UP026,
             // pydocstyle
             CheckKind::BlankLineAfterLastSection(..) => &CheckCode::D413,
             CheckKind::BlankLineAfterSection(..) => &CheckCode::D410,
@@ -2727,7 +2733,9 @@ 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(),
+            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 => {
@@ -3196,7 +3204,7 @@ impl CheckKind {
                 | CheckKind::ReplaceStdoutStderr
                 | CheckKind::ReplaceUniversalNewlines
                 | CheckKind::RewriteCElementTree
-                | CheckKind::RewriteMockImport
+                | CheckKind::RewriteMockImport(..)
                 | CheckKind::RewriteUnicodeLiteral
                 | CheckKind::SectionNameEndsInColon(..)
                 | CheckKind::SectionNotOverIndented(..)
@@ -3309,7 +3317,10 @@ 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::RewriteMockImport(reference_type) => Some(match reference_type {
+                MockReference::Import => "Import from `unittest.mock` instead".to_string(),
+                MockReference::Attribute => "Replace `mock.mock` with `mock`".to_string(),
+            }),
             CheckKind::NewLineAfterSectionName(name) => {
                 Some(format!("Add newline after \"{name}\""))
             }

diff --git a/src/pyupgrade/plugins/mod.rs b/src/pyupgrade/plugins/mod.rs
@@ -10,7 +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_mock_import::{rewrite_mock_attribute, 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;

diff --git a/src/pyupgrade/plugins/rewrite_mock_import.rs b/src/pyupgrade/plugins/rewrite_mock_import.rs
@@ -4,13 +4,14 @@ use libcst_native::{
     ImportAlias, ImportFrom, ImportNames, Name, NameOrAttribute, ParenthesizableWhitespace,
 };
 use log::error;
-use rustpython_ast::{Stmt, StmtKind};
+use rustpython_ast::{Expr, ExprKind, Stmt, StmtKind};
 
+use crate::ast::helpers::collect_call_paths;
 use crate::ast::types::Range;
 use crate::ast::whitespace::indentation;
 use crate::autofix::Fix;
 use crate::checkers::ast::Checker;
-use crate::checks::{Check, CheckCode, CheckKind};
+use crate::checks::{Check, CheckCode, CheckKind, MockReference};
 use crate::cst::matchers::{match_import, match_import_from, match_module};
 use crate::source_code_locator::SourceCodeLocator;
 use crate::source_code_style::SourceCodeStyleDetector;
@@ -177,6 +178,26 @@ fn format_import_from(
     })
 }
 
+/// UP026
+pub fn rewrite_mock_attribute(checker: &mut Checker, expr: &Expr) {
+    if let ExprKind::Attribute { value, .. } = &expr.node {
+        if collect_call_paths(value) == ["mock", "mock"] {
+            let mut check = Check::new(
+                CheckKind::RewriteMockImport(MockReference::Attribute),
+                Range::from_located(value),
+            );
+            if checker.patch(&CheckCode::UP026) {
+                check.amend(Fix::replacement(
+                    "mock".to_string(),
+                    value.location,
+                    value.end_location.unwrap(),
+                ));
+            }
+            checker.add_check(check);
+        }
+    }
+}
+
 /// UP026
 pub fn rewrite_mock_import(checker: &mut Checker, stmt: &Stmt) {
     match &stmt.node {
@@ -203,8 +224,10 @@ pub fn rewrite_mock_import(checker: &mut Checker, stmt: &Stmt) {
                 // Add a `Check` for each `mock` import.
                 for name in names {
                     if name.node.name == "mock" || name.node.name == "mock.mock" {
-                        let mut check =
-                            Check::new(CheckKind::RewriteMockImport, Range::from_located(name));
+                        let mut check = Check::new(
+                            CheckKind::RewriteMockImport(MockReference::Import),
+                            Range::from_located(name),
+                        );
                         if let Some(content) = content.as_ref() {
                             check.amend(Fix::replacement(
                                 content.clone(),
@@ -227,7 +250,10 @@ pub fn rewrite_mock_import(checker: &mut Checker, stmt: &Stmt) {
             }
 
             if module == "mock" {
-                let mut check = Check::new(CheckKind::RewriteMockImport, Range::from_located(stmt));
+                let mut check = Check::new(
+                    CheckKind::RewriteMockImport(MockReference::Import),
+                    Range::from_located(stmt),
+                );
                 if checker.patch(&CheckCode::UP026) {
                     let indent = indentation(checker, stmt);
                     match format_import_from(stmt, &indent, checker.locator, checker.style) {

diff --git a/src/pyupgrade/snapshots/ruff__pyupgrade__tests__UP026_UP026.py.snap b/src/pyupgrade/snapshots/ruff__pyupgrade__tests__UP026_UP026.py.snap
@@ -2,7 +2,8 @@
 source: src/pyupgrade/mod.rs
 expression: checks
 ---
-- kind: RewriteMockImport
+- kind:
+    RewriteMockImport: Import
   location:
     row: 3
     column: 11
@@ -18,7 +19,8 @@ expression: checks
       row: 3
       column: 15
   parent: ~
-- kind: RewriteMockImport
+- kind:
+    RewriteMockImport: Import
   location:
     row: 6
     column: 11
@@ -34,7 +36,8 @@ expression: checks
       row: 6
       column: 20
   parent: ~
-- kind: RewriteMockImport
+- kind:
+    RewriteMockImport: Import
   location:
     row: 9
     column: 7
@@ -50,7 +53,8 @@ expression: checks
       row: 9
       column: 16
   parent: ~
-- kind: RewriteMockImport
+- kind:
+    RewriteMockImport: Import
   location:
     row: 12
     column: 19
@@ -66,7 +70,8 @@ expression: checks
       row: 12
       column: 28
   parent: ~
-- kind: RewriteMockImport
+- kind:
+    RewriteMockImport: Import
   location:
     row: 15
     column: 7
@@ -82,7 +87,8 @@ expression: checks
       row: 15
       column: 16
   parent: ~
-- kind: RewriteMockImport
+- kind:
+    RewriteMockImport: Import
   location:
     row: 19
     column: 0
@@ -98,7 +104,8 @@ expression: checks
       row: 19
       column: 21
   parent: ~
-- kind: RewriteMockImport
+- kind:
+    RewriteMockImport: Import
   location:
     row: 22
     column: 0
@@ -114,7 +121,8 @@ expression: checks
       row: 27
       column: 1
   parent: ~
-- kind: RewriteMockImport
+- kind:
+    RewriteMockImport: Import
   location:
     row: 30
     column: 0
@@ -130,7 +138,8 @@ expression: checks
       row: 35
       column: 1
   parent: ~
-- kind: RewriteMockImport
+- kind:
+    RewriteMockImport: Import
   location:
     row: 39
     column: 8
@@ -146,7 +155,8 @@ expression: checks
       row: 44
       column: 9
   parent: ~
-- kind: RewriteMockImport
+- kind:
+    RewriteMockImport: Import
   location:
     row: 50
     column: 7
@@ -162,7 +172,8 @@ expression: checks
       row: 50
       column: 17
   parent: ~
-- kind: RewriteMockImport
+- kind:
+    RewriteMockImport: Import
   location:
     row: 50
     column: 13
@@ -178,7 +189,8 @@ expression: checks
       row: 50
       column: 17
   parent: ~
-- kind: RewriteMockImport
+- kind:
+    RewriteMockImport: Import
   location:
     row: 53
     column: 7
@@ -194,7 +206,8 @@ expression: checks
       row: 53
       column: 18
   parent: ~
-- kind: RewriteMockImport
+- kind:
+    RewriteMockImport: Import
   location:
     row: 56
     column: 0
@@ -210,7 +223,8 @@ expression: checks
       row: 56
       column: 28
   parent: ~
-- kind: RewriteMockImport
+- kind:
+    RewriteMockImport: Import
   location:
     row: 60
     column: 11
@@ -226,7 +240,8 @@ expression: checks
       row: 60
       column: 41
   parent: ~
-- kind: RewriteMockImport
+- kind:
+    RewriteMockImport: Import
   location:
     row: 60
     column: 24
@@ -242,7 +257,8 @@ expression: checks
       row: 60
       column: 41
   parent: ~
-- kind: RewriteMockImport
+- kind:
+    RewriteMockImport: Import
   location:
     row: 60
     column: 37
@@ -258,7 +274,8 @@ expression: checks
       row: 60
       column: 41
   parent: ~
-- kind: RewriteMockImport
+- kind:
+    RewriteMockImport: Import
   location:
     row: 63
     column: 11
@@ -274,7 +291,8 @@ expression: checks
       row: 63
       column: 45
   parent: ~
-- kind: RewriteMockImport
+- kind:
+    RewriteMockImport: Import
   location:
     row: 63
     column: 24
@@ -290,7 +308,8 @@ expression: checks
       row: 63
       column: 45
   parent: ~
-- kind: RewriteMockImport
+- kind:
+    RewriteMockImport: Import
   location:
     row: 63
     column: 37
@@ -306,7 +325,8 @@ expression: checks
       row: 63
       column: 45
   parent: ~
-- kind: RewriteMockImport
+- kind:
+    RewriteMockImport: Import
   location:
     row: 67
     column: 4
@@ -322,4 +342,21 @@ expression: checks
       row: 67
       column: 51
   parent: ~
+- kind:
+    RewriteMockImport: Attribute
+  location:
+    row: 74
+    column: 4
+  end_location:
+    row: 74
+    column: 13
+  fix:
+    content: mock
+    location:
+      row: 74
+      column: 4
+    end_location:
+      row: 74
+      column: 13
+  parent: ~