diff --git a/foo.py b/foo.py new file mode 100644 index 0000000000..4817a150a0 --- /dev/null +++ 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 index 15a369870e..e27774deb5 100644 --- a/resources/test/fixtures/pyupgrade/UP026.py +++ b/resources/test/fixtures/pyupgrade/UP026.py @@ -65,3 +65,10 @@ if True: 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 index 5672682195..83d17bdd30 100644 --- 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 index d6f012dbf2..5774c962ac 100644 --- a/src/checks.rs +++ b/src/checks.rs @@ -705,6 +705,12 @@ pub struct UnusedCodes { pub unmatched: Vec, } +#[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), 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 index 9033d1fadf..184444a147 100644 --- 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 index e300fc6a37..82ac7a2d98 100644 --- 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 index fd7b23507a..40b637eefa 100644 --- 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: ~