mirror of
https://github.com/astral-sh/ruff.git
synced 2025-10-09 18:02:19 +00:00
Pyupgrade: import mock
to from unittest import mock
(#1488)
This commit is contained in:
parent
f2c9f94f73
commit
70895a8f1e
14 changed files with 710 additions and 4 deletions
|
@ -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` | 🛠 |
|
| UP023 | RewriteCElementTree | `cElementTree` is deprecated, use `ElementTree` | 🛠 |
|
||||||
| UP024 | OSErrorAlias | Replace aliased errors with `OSError` | 🛠 |
|
| UP024 | OSErrorAlias | Replace aliased errors with `OSError` | 🛠 |
|
||||||
| UP025 | RewriteUnicodeLiteral | Remove unicode literals from strings | 🛠 |
|
| UP025 | RewriteUnicodeLiteral | Remove unicode literals from strings | 🛠 |
|
||||||
|
| UP026 | RewriteMockImport | `mock` is deprecated, use `unittest.mock` | 🛠 |
|
||||||
|
|
||||||
### pep8-naming (N)
|
### pep8-naming (N)
|
||||||
|
|
||||||
|
|
13
foo.py
Normal file
13
foo.py
Normal file
|
@ -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
|
67
resources/test/fixtures/pyupgrade/UP026.py
vendored
Normal file
67
resources/test/fixtures/pyupgrade/UP026.py
vendored
Normal file
|
@ -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
|
|
@ -911,6 +911,7 @@
|
||||||
"UP023",
|
"UP023",
|
||||||
"UP024",
|
"UP024",
|
||||||
"UP025",
|
"UP025",
|
||||||
|
"UP026",
|
||||||
"W",
|
"W",
|
||||||
"W2",
|
"W2",
|
||||||
"W29",
|
"W29",
|
||||||
|
|
|
@ -643,6 +643,9 @@ where
|
||||||
if self.settings.enabled.contains(&CheckCode::UP023) {
|
if self.settings.enabled.contains(&CheckCode::UP023) {
|
||||||
pyupgrade::plugins::replace_c_element_tree(self, stmt);
|
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 {
|
for alias in names {
|
||||||
if alias.node.name.contains('.') && alias.node.asname.is_none() {
|
if alias.node.name.contains('.') && alias.node.asname.is_none() {
|
||||||
|
@ -854,6 +857,9 @@ where
|
||||||
pyupgrade::plugins::unnecessary_future_import(self, stmt, names);
|
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 self.settings.enabled.contains(&CheckCode::TID251) {
|
||||||
if let Some(module) = module {
|
if let Some(module) = module {
|
||||||
|
|
|
@ -238,6 +238,7 @@ pub enum CheckCode {
|
||||||
UP023,
|
UP023,
|
||||||
UP024,
|
UP024,
|
||||||
UP025,
|
UP025,
|
||||||
|
UP026,
|
||||||
// pydocstyle
|
// pydocstyle
|
||||||
D100,
|
D100,
|
||||||
D101,
|
D101,
|
||||||
|
@ -910,6 +911,7 @@ pub enum CheckKind {
|
||||||
RewriteCElementTree,
|
RewriteCElementTree,
|
||||||
OSErrorAlias(Option<String>),
|
OSErrorAlias(Option<String>),
|
||||||
RewriteUnicodeLiteral,
|
RewriteUnicodeLiteral,
|
||||||
|
RewriteMockImport,
|
||||||
// pydocstyle
|
// pydocstyle
|
||||||
BlankLineAfterLastSection(String),
|
BlankLineAfterLastSection(String),
|
||||||
BlankLineAfterSection(String),
|
BlankLineAfterSection(String),
|
||||||
|
@ -1305,6 +1307,7 @@ impl CheckCode {
|
||||||
CheckCode::UP023 => CheckKind::RewriteCElementTree,
|
CheckCode::UP023 => CheckKind::RewriteCElementTree,
|
||||||
CheckCode::UP024 => CheckKind::OSErrorAlias(None),
|
CheckCode::UP024 => CheckKind::OSErrorAlias(None),
|
||||||
CheckCode::UP025 => CheckKind::RewriteUnicodeLiteral,
|
CheckCode::UP025 => CheckKind::RewriteUnicodeLiteral,
|
||||||
|
CheckCode::UP026 => CheckKind::RewriteMockImport,
|
||||||
// pydocstyle
|
// pydocstyle
|
||||||
CheckCode::D100 => CheckKind::PublicModule,
|
CheckCode::D100 => CheckKind::PublicModule,
|
||||||
CheckCode::D101 => CheckKind::PublicClass,
|
CheckCode::D101 => CheckKind::PublicClass,
|
||||||
|
@ -1738,6 +1741,7 @@ impl CheckCode {
|
||||||
CheckCode::UP023 => CheckCategory::Pyupgrade,
|
CheckCode::UP023 => CheckCategory::Pyupgrade,
|
||||||
CheckCode::UP024 => CheckCategory::Pyupgrade,
|
CheckCode::UP024 => CheckCategory::Pyupgrade,
|
||||||
CheckCode::UP025 => CheckCategory::Pyupgrade,
|
CheckCode::UP025 => CheckCategory::Pyupgrade,
|
||||||
|
CheckCode::UP026 => CheckCategory::Pyupgrade,
|
||||||
CheckCode::W292 => CheckCategory::Pycodestyle,
|
CheckCode::W292 => CheckCategory::Pycodestyle,
|
||||||
CheckCode::W605 => CheckCategory::Pycodestyle,
|
CheckCode::W605 => CheckCategory::Pycodestyle,
|
||||||
CheckCode::YTT101 => CheckCategory::Flake82020,
|
CheckCode::YTT101 => CheckCategory::Flake82020,
|
||||||
|
@ -1961,6 +1965,7 @@ impl CheckKind {
|
||||||
CheckKind::RewriteCElementTree => &CheckCode::UP023,
|
CheckKind::RewriteCElementTree => &CheckCode::UP023,
|
||||||
CheckKind::OSErrorAlias(..) => &CheckCode::UP024,
|
CheckKind::OSErrorAlias(..) => &CheckCode::UP024,
|
||||||
CheckKind::RewriteUnicodeLiteral => &CheckCode::UP025,
|
CheckKind::RewriteUnicodeLiteral => &CheckCode::UP025,
|
||||||
|
CheckKind::RewriteMockImport => &CheckCode::UP026,
|
||||||
// pydocstyle
|
// pydocstyle
|
||||||
CheckKind::BlankLineAfterLastSection(..) => &CheckCode::D413,
|
CheckKind::BlankLineAfterLastSection(..) => &CheckCode::D413,
|
||||||
CheckKind::BlankLineAfterSection(..) => &CheckCode::D410,
|
CheckKind::BlankLineAfterSection(..) => &CheckCode::D410,
|
||||||
|
@ -2722,6 +2727,7 @@ impl CheckKind {
|
||||||
}
|
}
|
||||||
CheckKind::OSErrorAlias(..) => "Replace aliased errors with `OSError`".to_string(),
|
CheckKind::OSErrorAlias(..) => "Replace aliased errors with `OSError`".to_string(),
|
||||||
CheckKind::RewriteUnicodeLiteral => "Remove unicode literals from strings".to_string(),
|
CheckKind::RewriteUnicodeLiteral => "Remove unicode literals from strings".to_string(),
|
||||||
|
CheckKind::RewriteMockImport => "`mock` is deprecated, use `unittest.mock`".to_string(),
|
||||||
// pydocstyle
|
// pydocstyle
|
||||||
CheckKind::FitsOnOneLine => "One-line docstring should fit on one line".to_string(),
|
CheckKind::FitsOnOneLine => "One-line docstring should fit on one line".to_string(),
|
||||||
CheckKind::BlankLineAfterSummary => {
|
CheckKind::BlankLineAfterSummary => {
|
||||||
|
@ -3190,6 +3196,7 @@ impl CheckKind {
|
||||||
| CheckKind::ReplaceStdoutStderr
|
| CheckKind::ReplaceStdoutStderr
|
||||||
| CheckKind::ReplaceUniversalNewlines
|
| CheckKind::ReplaceUniversalNewlines
|
||||||
| CheckKind::RewriteCElementTree
|
| CheckKind::RewriteCElementTree
|
||||||
|
| CheckKind::RewriteMockImport
|
||||||
| CheckKind::RewriteUnicodeLiteral
|
| CheckKind::RewriteUnicodeLiteral
|
||||||
| CheckKind::SectionNameEndsInColon(..)
|
| CheckKind::SectionNameEndsInColon(..)
|
||||||
| CheckKind::SectionNotOverIndented(..)
|
| CheckKind::SectionNotOverIndented(..)
|
||||||
|
@ -3302,6 +3309,7 @@ impl CheckKind {
|
||||||
}
|
}
|
||||||
CheckKind::RewriteCElementTree => Some("Replace with `ElementTree`".to_string()),
|
CheckKind::RewriteCElementTree => Some("Replace with `ElementTree`".to_string()),
|
||||||
CheckKind::RewriteUnicodeLiteral => Some("Remove unicode prefix".to_string()),
|
CheckKind::RewriteUnicodeLiteral => Some("Remove unicode prefix".to_string()),
|
||||||
|
CheckKind::RewriteMockImport => Some("Import from `unittest.mock` instead".to_string()),
|
||||||
CheckKind::NewLineAfterSectionName(name) => {
|
CheckKind::NewLineAfterSectionName(name) => {
|
||||||
Some(format!("Add newline after \"{name}\""))
|
Some(format!("Add newline after \"{name}\""))
|
||||||
}
|
}
|
||||||
|
|
|
@ -542,6 +542,7 @@ pub enum CheckCodePrefix {
|
||||||
UP023,
|
UP023,
|
||||||
UP024,
|
UP024,
|
||||||
UP025,
|
UP025,
|
||||||
|
UP026,
|
||||||
W,
|
W,
|
||||||
W2,
|
W2,
|
||||||
W29,
|
W29,
|
||||||
|
@ -776,6 +777,7 @@ impl CheckCodePrefix {
|
||||||
CheckCode::UP023,
|
CheckCode::UP023,
|
||||||
CheckCode::UP024,
|
CheckCode::UP024,
|
||||||
CheckCode::UP025,
|
CheckCode::UP025,
|
||||||
|
CheckCode::UP026,
|
||||||
CheckCode::D100,
|
CheckCode::D100,
|
||||||
CheckCode::D101,
|
CheckCode::D101,
|
||||||
CheckCode::D102,
|
CheckCode::D102,
|
||||||
|
@ -2456,6 +2458,7 @@ impl CheckCodePrefix {
|
||||||
CheckCode::UP023,
|
CheckCode::UP023,
|
||||||
CheckCode::UP024,
|
CheckCode::UP024,
|
||||||
CheckCode::UP025,
|
CheckCode::UP025,
|
||||||
|
CheckCode::UP026,
|
||||||
]
|
]
|
||||||
}
|
}
|
||||||
CheckCodePrefix::U0 => {
|
CheckCodePrefix::U0 => {
|
||||||
|
@ -2490,6 +2493,7 @@ impl CheckCodePrefix {
|
||||||
CheckCode::UP023,
|
CheckCode::UP023,
|
||||||
CheckCode::UP024,
|
CheckCode::UP024,
|
||||||
CheckCode::UP025,
|
CheckCode::UP025,
|
||||||
|
CheckCode::UP026,
|
||||||
]
|
]
|
||||||
}
|
}
|
||||||
CheckCodePrefix::U00 => {
|
CheckCodePrefix::U00 => {
|
||||||
|
@ -2708,6 +2712,7 @@ impl CheckCodePrefix {
|
||||||
CheckCode::UP023,
|
CheckCode::UP023,
|
||||||
CheckCode::UP024,
|
CheckCode::UP024,
|
||||||
CheckCode::UP025,
|
CheckCode::UP025,
|
||||||
|
CheckCode::UP026,
|
||||||
],
|
],
|
||||||
CheckCodePrefix::UP0 => vec![
|
CheckCodePrefix::UP0 => vec![
|
||||||
CheckCode::UP001,
|
CheckCode::UP001,
|
||||||
|
@ -2734,6 +2739,7 @@ impl CheckCodePrefix {
|
||||||
CheckCode::UP023,
|
CheckCode::UP023,
|
||||||
CheckCode::UP024,
|
CheckCode::UP024,
|
||||||
CheckCode::UP025,
|
CheckCode::UP025,
|
||||||
|
CheckCode::UP026,
|
||||||
],
|
],
|
||||||
CheckCodePrefix::UP00 => vec![
|
CheckCodePrefix::UP00 => vec![
|
||||||
CheckCode::UP001,
|
CheckCode::UP001,
|
||||||
|
@ -2782,6 +2788,7 @@ impl CheckCodePrefix {
|
||||||
CheckCode::UP023,
|
CheckCode::UP023,
|
||||||
CheckCode::UP024,
|
CheckCode::UP024,
|
||||||
CheckCode::UP025,
|
CheckCode::UP025,
|
||||||
|
CheckCode::UP026,
|
||||||
],
|
],
|
||||||
CheckCodePrefix::UP020 => vec![CheckCode::UP020],
|
CheckCodePrefix::UP020 => vec![CheckCode::UP020],
|
||||||
CheckCodePrefix::UP021 => vec![CheckCode::UP021],
|
CheckCodePrefix::UP021 => vec![CheckCode::UP021],
|
||||||
|
@ -2789,6 +2796,7 @@ impl CheckCodePrefix {
|
||||||
CheckCodePrefix::UP023 => vec![CheckCode::UP023],
|
CheckCodePrefix::UP023 => vec![CheckCode::UP023],
|
||||||
CheckCodePrefix::UP024 => vec![CheckCode::UP024],
|
CheckCodePrefix::UP024 => vec![CheckCode::UP024],
|
||||||
CheckCodePrefix::UP025 => vec![CheckCode::UP025],
|
CheckCodePrefix::UP025 => vec![CheckCode::UP025],
|
||||||
|
CheckCodePrefix::UP026 => vec![CheckCode::UP026],
|
||||||
CheckCodePrefix::W => vec![CheckCode::W292, CheckCode::W605],
|
CheckCodePrefix::W => vec![CheckCode::W292, CheckCode::W605],
|
||||||
CheckCodePrefix::W2 => vec![CheckCode::W292],
|
CheckCodePrefix::W2 => vec![CheckCode::W292],
|
||||||
CheckCodePrefix::W29 => vec![CheckCode::W292],
|
CheckCodePrefix::W29 => vec![CheckCode::W292],
|
||||||
|
@ -3362,6 +3370,7 @@ impl CheckCodePrefix {
|
||||||
CheckCodePrefix::UP023 => SuffixLength::Three,
|
CheckCodePrefix::UP023 => SuffixLength::Three,
|
||||||
CheckCodePrefix::UP024 => SuffixLength::Three,
|
CheckCodePrefix::UP024 => SuffixLength::Three,
|
||||||
CheckCodePrefix::UP025 => SuffixLength::Three,
|
CheckCodePrefix::UP025 => SuffixLength::Three,
|
||||||
|
CheckCodePrefix::UP026 => SuffixLength::Three,
|
||||||
CheckCodePrefix::W => SuffixLength::Zero,
|
CheckCodePrefix::W => SuffixLength::Zero,
|
||||||
CheckCodePrefix::W2 => SuffixLength::One,
|
CheckCodePrefix::W2 => SuffixLength::One,
|
||||||
CheckCodePrefix::W29 => SuffixLength::Two,
|
CheckCodePrefix::W29 => SuffixLength::Two,
|
||||||
|
|
|
@ -1,5 +1,5 @@
|
||||||
use anyhow::{bail, Result};
|
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> {
|
pub fn match_module(module_text: &str) -> Result<Module> {
|
||||||
match libcst_native::parse_module(module_text, None) {
|
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")
|
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")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
|
@ -23,12 +23,12 @@ use crate::SourceCodeLocator;
|
||||||
mod categorize;
|
mod categorize;
|
||||||
mod comments;
|
mod comments;
|
||||||
pub mod format;
|
pub mod format;
|
||||||
mod helpers;
|
pub mod helpers;
|
||||||
pub mod plugins;
|
pub mod plugins;
|
||||||
pub mod settings;
|
pub mod settings;
|
||||||
mod sorting;
|
mod sorting;
|
||||||
pub mod track;
|
pub mod track;
|
||||||
mod types;
|
pub mod types;
|
||||||
|
|
||||||
#[derive(Debug)]
|
#[derive(Debug)]
|
||||||
pub struct AnnotatedAliasData<'a> {
|
pub struct AnnotatedAliasData<'a> {
|
||||||
|
|
|
@ -46,6 +46,7 @@ mod tests {
|
||||||
#[test_case(CheckCode::UP024, Path::new("UP024_1.py"); "UP024_1")]
|
#[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::UP024, Path::new("UP024_2.py"); "UP024_2")]
|
||||||
#[test_case(CheckCode::UP025, Path::new("UP025.py"); "UP025")]
|
#[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<()> {
|
fn checks(check_code: CheckCode, path: &Path) -> Result<()> {
|
||||||
let snapshot = format!("{}_{}", check_code.as_ref(), path.to_string_lossy());
|
let snapshot = format!("{}_{}", check_code.as_ref(), path.to_string_lossy());
|
||||||
let checks = test_path(
|
let checks = test_path(
|
||||||
|
|
|
@ -10,6 +10,7 @@ pub use remove_six_compat::remove_six_compat;
|
||||||
pub use replace_stdout_stderr::replace_stdout_stderr;
|
pub use replace_stdout_stderr::replace_stdout_stderr;
|
||||||
pub use replace_universal_newlines::replace_universal_newlines;
|
pub use replace_universal_newlines::replace_universal_newlines;
|
||||||
pub use rewrite_c_element_tree::replace_c_element_tree;
|
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 rewrite_unicode_literal::rewrite_unicode_literal;
|
||||||
pub use super_call_with_parameters::super_call_with_parameters;
|
pub use super_call_with_parameters::super_call_with_parameters;
|
||||||
pub use type_of_primitive::type_of_primitive;
|
pub use type_of_primitive::type_of_primitive;
|
||||||
|
@ -34,6 +35,7 @@ mod remove_six_compat;
|
||||||
mod replace_stdout_stderr;
|
mod replace_stdout_stderr;
|
||||||
mod replace_universal_newlines;
|
mod replace_universal_newlines;
|
||||||
mod rewrite_c_element_tree;
|
mod rewrite_c_element_tree;
|
||||||
|
mod rewrite_mock_import;
|
||||||
mod rewrite_unicode_literal;
|
mod rewrite_unicode_literal;
|
||||||
mod super_call_with_parameters;
|
mod super_call_with_parameters;
|
||||||
mod type_of_primitive;
|
mod type_of_primitive;
|
||||||
|
|
249
src/pyupgrade/plugins/rewrite_mock_import.rs
Normal file
249
src/pyupgrade/plugins/rewrite_mock_import.rs
Normal file
|
@ -0,0 +1,249 @@
|
||||||
|
use anyhow::Result;
|
||||||
|
use libcst_native::{
|
||||||
|
AsName, AssignTargetExpression, Attribute, Codegen, CodegenState, Dot, Expression, Import,
|
||||||
|
ImportAlias, ImportFrom, ImportNames, Name, NameOrAttribute, ParenthesizableWhitespace,
|
||||||
|
};
|
||||||
|
use log::error;
|
||||||
|
use rustpython_ast::{Stmt, StmtKind};
|
||||||
|
|
||||||
|
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::cst::matchers::{match_import, match_import_from, match_module};
|
||||||
|
use crate::source_code_locator::SourceCodeLocator;
|
||||||
|
use crate::source_code_style::SourceCodeStyleDetector;
|
||||||
|
|
||||||
|
/// Return a vector of all non-`mock` imports.
|
||||||
|
fn clean_import_aliases(aliases: Vec<ImportAlias>) -> (Vec<ImportAlias>, Vec<Option<AsName>>) {
|
||||||
|
let mut clean_aliases: Vec<ImportAlias> = vec![];
|
||||||
|
let mut mock_aliases: Vec<Option<AsName>> = vec![];
|
||||||
|
for alias in aliases {
|
||||||
|
match &alias.name {
|
||||||
|
// Ex) `import mock`
|
||||||
|
NameOrAttribute::N(name_struct) => {
|
||||||
|
if name_struct.value == "mock" {
|
||||||
|
mock_aliases.push(alias.asname.clone());
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
clean_aliases.push(alias);
|
||||||
|
}
|
||||||
|
// Ex) `import mock.mock`
|
||||||
|
NameOrAttribute::A(attribute_struct) => {
|
||||||
|
if let Expression::Name(name_struct) = &*attribute_struct.value {
|
||||||
|
if name_struct.value == "mock" && attribute_struct.attr.value == "mock" {
|
||||||
|
mock_aliases.push(alias.asname.clone());
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
clean_aliases.push(alias);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
(clean_aliases, mock_aliases)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Return `true` if the aliases contain `mock`.
|
||||||
|
fn includes_mock_member(aliases: &[ImportAlias]) -> bool {
|
||||||
|
for alias in aliases {
|
||||||
|
let ImportAlias { name, .. } = &alias;
|
||||||
|
// Ex) `import mock.mock`
|
||||||
|
if let NameOrAttribute::A(attribute_struct) = name {
|
||||||
|
if let Expression::Name(name_struct) = &*attribute_struct.value {
|
||||||
|
if name_struct.value == "mock" && attribute_struct.attr.value == "mock" {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
false
|
||||||
|
}
|
||||||
|
|
||||||
|
fn format_mocks(
|
||||||
|
aliases: Vec<Option<AsName>>,
|
||||||
|
indent: &str,
|
||||||
|
stylist: &SourceCodeStyleDetector,
|
||||||
|
) -> String {
|
||||||
|
let mut content = String::new();
|
||||||
|
for alias in aliases {
|
||||||
|
match alias {
|
||||||
|
None => {
|
||||||
|
if !content.is_empty() {
|
||||||
|
content.push_str(stylist.line_ending());
|
||||||
|
content.push_str(indent);
|
||||||
|
}
|
||||||
|
content.push_str("from unittest import mock");
|
||||||
|
}
|
||||||
|
Some(as_name) => {
|
||||||
|
if let AssignTargetExpression::Name(name) = as_name.name {
|
||||||
|
if !content.is_empty() {
|
||||||
|
content.push_str(stylist.line_ending());
|
||||||
|
content.push_str(indent);
|
||||||
|
}
|
||||||
|
content.push_str("from unittest import mock as ");
|
||||||
|
content.push_str(name.value);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
content
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Format the `import mock` rewrite.
|
||||||
|
fn format_import(
|
||||||
|
stmt: &Stmt,
|
||||||
|
indent: &str,
|
||||||
|
locator: &SourceCodeLocator,
|
||||||
|
stylist: &SourceCodeStyleDetector,
|
||||||
|
) -> Result<String> {
|
||||||
|
let module_text = locator.slice_source_code_range(&Range::from_located(stmt));
|
||||||
|
let mut tree = match_module(&module_text)?;
|
||||||
|
let mut import = match_import(&mut tree)?;
|
||||||
|
|
||||||
|
let Import { names, .. } = import.clone();
|
||||||
|
let (clean_aliases, mock_aliases) = clean_import_aliases(names);
|
||||||
|
|
||||||
|
Ok(if clean_aliases.is_empty() {
|
||||||
|
format_mocks(mock_aliases, indent, stylist)
|
||||||
|
} else {
|
||||||
|
import.names = clean_aliases;
|
||||||
|
|
||||||
|
let mut state = CodegenState::default();
|
||||||
|
tree.codegen(&mut state);
|
||||||
|
|
||||||
|
let mut content = state.to_string();
|
||||||
|
content.push_str(stylist.line_ending());
|
||||||
|
content.push_str(indent);
|
||||||
|
content.push_str(&format_mocks(mock_aliases, indent, stylist));
|
||||||
|
content
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Format the `from mock import ...` rewrite.
|
||||||
|
fn format_import_from(
|
||||||
|
stmt: &Stmt,
|
||||||
|
indent: &str,
|
||||||
|
locator: &SourceCodeLocator,
|
||||||
|
stylist: &SourceCodeStyleDetector,
|
||||||
|
) -> Result<String> {
|
||||||
|
let module_text = locator.slice_source_code_range(&Range::from_located(stmt));
|
||||||
|
let mut tree = match_module(&module_text).unwrap();
|
||||||
|
let mut import = match_import_from(&mut tree)?;
|
||||||
|
|
||||||
|
let ImportFrom {
|
||||||
|
names: ImportNames::Aliases(names),
|
||||||
|
..
|
||||||
|
} = import.clone() else {
|
||||||
|
unreachable!("Expected ImportNames::Aliases");
|
||||||
|
};
|
||||||
|
|
||||||
|
let has_mock_member = includes_mock_member(&names);
|
||||||
|
let (clean_aliases, mock_aliases) = clean_import_aliases(names);
|
||||||
|
|
||||||
|
Ok(if clean_aliases.is_empty() {
|
||||||
|
format_mocks(mock_aliases, indent, stylist)
|
||||||
|
} else {
|
||||||
|
import.names = ImportNames::Aliases(clean_aliases);
|
||||||
|
import.module = Some(NameOrAttribute::A(Box::new(Attribute {
|
||||||
|
value: Box::new(Expression::Name(Box::new(Name {
|
||||||
|
value: "unittest",
|
||||||
|
lpar: vec![],
|
||||||
|
rpar: vec![],
|
||||||
|
}))),
|
||||||
|
attr: Name {
|
||||||
|
value: "mock",
|
||||||
|
lpar: vec![],
|
||||||
|
rpar: vec![],
|
||||||
|
},
|
||||||
|
dot: Dot {
|
||||||
|
whitespace_before: ParenthesizableWhitespace::default(),
|
||||||
|
whitespace_after: ParenthesizableWhitespace::default(),
|
||||||
|
},
|
||||||
|
lpar: vec![],
|
||||||
|
rpar: vec![],
|
||||||
|
})));
|
||||||
|
|
||||||
|
let mut state = CodegenState::default();
|
||||||
|
tree.codegen(&mut state);
|
||||||
|
|
||||||
|
let mut content = state.to_string();
|
||||||
|
if has_mock_member {
|
||||||
|
content.push_str(stylist.line_ending());
|
||||||
|
content.push_str(indent);
|
||||||
|
content.push_str(&format_mocks(mock_aliases, indent, stylist));
|
||||||
|
}
|
||||||
|
content
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// UP026
|
||||||
|
pub fn rewrite_mock_import(checker: &mut Checker, stmt: &Stmt) {
|
||||||
|
match &stmt.node {
|
||||||
|
StmtKind::Import { names } => {
|
||||||
|
// Find all `mock` imports.
|
||||||
|
if names
|
||||||
|
.iter()
|
||||||
|
.any(|name| name.node.name == "mock" || name.node.name == "mock.mock")
|
||||||
|
{
|
||||||
|
// Generate the fix, if needed, which is shared between all `mock` imports.
|
||||||
|
let content = if checker.patch(&CheckCode::UP026) {
|
||||||
|
let indent = indentation(checker, stmt);
|
||||||
|
match format_import(stmt, &indent, checker.locator, checker.style) {
|
||||||
|
Ok(content) => Some(content),
|
||||||
|
Err(e) => {
|
||||||
|
error!("Failed to rewrite `mock` import: {e}");
|
||||||
|
None
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
|
||||||
|
// 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));
|
||||||
|
if let Some(content) = content.as_ref() {
|
||||||
|
check.amend(Fix::replacement(
|
||||||
|
content.clone(),
|
||||||
|
stmt.location,
|
||||||
|
stmt.end_location.unwrap(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
checker.add_check(check);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
StmtKind::ImportFrom {
|
||||||
|
module: Some(module),
|
||||||
|
level,
|
||||||
|
..
|
||||||
|
} => {
|
||||||
|
if level.map_or(false, |level| level > 0) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
if module == "mock" {
|
||||||
|
let mut check = Check::new(CheckKind::RewriteMockImport, Range::from_located(stmt));
|
||||||
|
if checker.patch(&CheckCode::UP026) {
|
||||||
|
let indent = indentation(checker, stmt);
|
||||||
|
match format_import_from(stmt, &indent, checker.locator, checker.style) {
|
||||||
|
Ok(content) => {
|
||||||
|
check.amend(Fix::replacement(
|
||||||
|
content,
|
||||||
|
stmt.location,
|
||||||
|
stmt.end_location.unwrap(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
Err(e) => error!("Failed to rewrite `mock` import: {e}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
checker.add_check(check);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => (),
|
||||||
|
}
|
||||||
|
}
|
|
@ -80,7 +80,7 @@ pub fn unnecessary_future_import(checker: &mut Checker, stmt: &Stmt, names: &[Lo
|
||||||
}
|
}
|
||||||
check.amend(fix);
|
check.amend(fix);
|
||||||
}
|
}
|
||||||
Err(e) => error!("Failed to remove __future__ import: {e}"),
|
Err(e) => error!("Failed to remove `__future__` import: {e}"),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
checker.add_check(check);
|
checker.add_check(check);
|
||||||
|
|
|
@ -0,0 +1,325 @@
|
||||||
|
---
|
||||||
|
source: src/pyupgrade/mod.rs
|
||||||
|
expression: checks
|
||||||
|
---
|
||||||
|
- kind: RewriteMockImport
|
||||||
|
location:
|
||||||
|
row: 3
|
||||||
|
column: 11
|
||||||
|
end_location:
|
||||||
|
row: 3
|
||||||
|
column: 15
|
||||||
|
fix:
|
||||||
|
content: from unittest import mock
|
||||||
|
location:
|
||||||
|
row: 3
|
||||||
|
column: 4
|
||||||
|
end_location:
|
||||||
|
row: 3
|
||||||
|
column: 15
|
||||||
|
parent: ~
|
||||||
|
- kind: RewriteMockImport
|
||||||
|
location:
|
||||||
|
row: 6
|
||||||
|
column: 11
|
||||||
|
end_location:
|
||||||
|
row: 6
|
||||||
|
column: 15
|
||||||
|
fix:
|
||||||
|
content: "import sys\n from unittest import mock"
|
||||||
|
location:
|
||||||
|
row: 6
|
||||||
|
column: 4
|
||||||
|
end_location:
|
||||||
|
row: 6
|
||||||
|
column: 20
|
||||||
|
parent: ~
|
||||||
|
- kind: RewriteMockImport
|
||||||
|
location:
|
||||||
|
row: 9
|
||||||
|
column: 7
|
||||||
|
end_location:
|
||||||
|
row: 9
|
||||||
|
column: 16
|
||||||
|
fix:
|
||||||
|
content: from unittest import mock
|
||||||
|
location:
|
||||||
|
row: 9
|
||||||
|
column: 0
|
||||||
|
end_location:
|
||||||
|
row: 9
|
||||||
|
column: 16
|
||||||
|
parent: ~
|
||||||
|
- kind: RewriteMockImport
|
||||||
|
location:
|
||||||
|
row: 12
|
||||||
|
column: 19
|
||||||
|
end_location:
|
||||||
|
row: 12
|
||||||
|
column: 23
|
||||||
|
fix:
|
||||||
|
content: "import contextlib, sys\nfrom unittest import mock"
|
||||||
|
location:
|
||||||
|
row: 12
|
||||||
|
column: 0
|
||||||
|
end_location:
|
||||||
|
row: 12
|
||||||
|
column: 28
|
||||||
|
parent: ~
|
||||||
|
- kind: RewriteMockImport
|
||||||
|
location:
|
||||||
|
row: 15
|
||||||
|
column: 7
|
||||||
|
end_location:
|
||||||
|
row: 15
|
||||||
|
column: 11
|
||||||
|
fix:
|
||||||
|
content: "import sys\nfrom unittest import mock"
|
||||||
|
location:
|
||||||
|
row: 15
|
||||||
|
column: 0
|
||||||
|
end_location:
|
||||||
|
row: 15
|
||||||
|
column: 16
|
||||||
|
parent: ~
|
||||||
|
- kind: RewriteMockImport
|
||||||
|
location:
|
||||||
|
row: 19
|
||||||
|
column: 0
|
||||||
|
end_location:
|
||||||
|
row: 19
|
||||||
|
column: 21
|
||||||
|
fix:
|
||||||
|
content: from unittest import mock
|
||||||
|
location:
|
||||||
|
row: 19
|
||||||
|
column: 0
|
||||||
|
end_location:
|
||||||
|
row: 19
|
||||||
|
column: 21
|
||||||
|
parent: ~
|
||||||
|
- kind: RewriteMockImport
|
||||||
|
location:
|
||||||
|
row: 22
|
||||||
|
column: 0
|
||||||
|
end_location:
|
||||||
|
row: 27
|
||||||
|
column: 1
|
||||||
|
fix:
|
||||||
|
content: "from unittest.mock import (\n a,\n b,\n c,\n)"
|
||||||
|
location:
|
||||||
|
row: 22
|
||||||
|
column: 0
|
||||||
|
end_location:
|
||||||
|
row: 27
|
||||||
|
column: 1
|
||||||
|
parent: ~
|
||||||
|
- kind: RewriteMockImport
|
||||||
|
location:
|
||||||
|
row: 30
|
||||||
|
column: 0
|
||||||
|
end_location:
|
||||||
|
row: 35
|
||||||
|
column: 1
|
||||||
|
fix:
|
||||||
|
content: "from unittest.mock import (\n a,\n b,\n c\n)"
|
||||||
|
location:
|
||||||
|
row: 30
|
||||||
|
column: 0
|
||||||
|
end_location:
|
||||||
|
row: 35
|
||||||
|
column: 1
|
||||||
|
parent: ~
|
||||||
|
- kind: RewriteMockImport
|
||||||
|
location:
|
||||||
|
row: 39
|
||||||
|
column: 8
|
||||||
|
end_location:
|
||||||
|
row: 44
|
||||||
|
column: 9
|
||||||
|
fix:
|
||||||
|
content: "from unittest.mock import (\n a,\n b,\n c\n )"
|
||||||
|
location:
|
||||||
|
row: 39
|
||||||
|
column: 8
|
||||||
|
end_location:
|
||||||
|
row: 44
|
||||||
|
column: 9
|
||||||
|
parent: ~
|
||||||
|
- kind: RewriteMockImport
|
||||||
|
location:
|
||||||
|
row: 50
|
||||||
|
column: 7
|
||||||
|
end_location:
|
||||||
|
row: 50
|
||||||
|
column: 11
|
||||||
|
fix:
|
||||||
|
content: "from unittest import mock\nfrom unittest import mock"
|
||||||
|
location:
|
||||||
|
row: 50
|
||||||
|
column: 0
|
||||||
|
end_location:
|
||||||
|
row: 50
|
||||||
|
column: 17
|
||||||
|
parent: ~
|
||||||
|
- kind: RewriteMockImport
|
||||||
|
location:
|
||||||
|
row: 50
|
||||||
|
column: 13
|
||||||
|
end_location:
|
||||||
|
row: 50
|
||||||
|
column: 17
|
||||||
|
fix:
|
||||||
|
content: "from unittest import mock\nfrom unittest import mock"
|
||||||
|
location:
|
||||||
|
row: 50
|
||||||
|
column: 0
|
||||||
|
end_location:
|
||||||
|
row: 50
|
||||||
|
column: 17
|
||||||
|
parent: ~
|
||||||
|
- kind: RewriteMockImport
|
||||||
|
location:
|
||||||
|
row: 53
|
||||||
|
column: 7
|
||||||
|
end_location:
|
||||||
|
row: 53
|
||||||
|
column: 18
|
||||||
|
fix:
|
||||||
|
content: from unittest import mock as foo
|
||||||
|
location:
|
||||||
|
row: 53
|
||||||
|
column: 0
|
||||||
|
end_location:
|
||||||
|
row: 53
|
||||||
|
column: 18
|
||||||
|
parent: ~
|
||||||
|
- kind: RewriteMockImport
|
||||||
|
location:
|
||||||
|
row: 56
|
||||||
|
column: 0
|
||||||
|
end_location:
|
||||||
|
row: 56
|
||||||
|
column: 28
|
||||||
|
fix:
|
||||||
|
content: from unittest import mock as foo
|
||||||
|
location:
|
||||||
|
row: 56
|
||||||
|
column: 0
|
||||||
|
end_location:
|
||||||
|
row: 56
|
||||||
|
column: 28
|
||||||
|
parent: ~
|
||||||
|
- kind: RewriteMockImport
|
||||||
|
location:
|
||||||
|
row: 60
|
||||||
|
column: 11
|
||||||
|
end_location:
|
||||||
|
row: 60
|
||||||
|
column: 22
|
||||||
|
fix:
|
||||||
|
content: "from unittest import mock as foo\n from unittest import mock as bar\n from unittest import mock"
|
||||||
|
location:
|
||||||
|
row: 60
|
||||||
|
column: 4
|
||||||
|
end_location:
|
||||||
|
row: 60
|
||||||
|
column: 41
|
||||||
|
parent: ~
|
||||||
|
- kind: RewriteMockImport
|
||||||
|
location:
|
||||||
|
row: 60
|
||||||
|
column: 24
|
||||||
|
end_location:
|
||||||
|
row: 60
|
||||||
|
column: 35
|
||||||
|
fix:
|
||||||
|
content: "from unittest import mock as foo\n from unittest import mock as bar\n from unittest import mock"
|
||||||
|
location:
|
||||||
|
row: 60
|
||||||
|
column: 4
|
||||||
|
end_location:
|
||||||
|
row: 60
|
||||||
|
column: 41
|
||||||
|
parent: ~
|
||||||
|
- kind: RewriteMockImport
|
||||||
|
location:
|
||||||
|
row: 60
|
||||||
|
column: 37
|
||||||
|
end_location:
|
||||||
|
row: 60
|
||||||
|
column: 41
|
||||||
|
fix:
|
||||||
|
content: "from unittest import mock as foo\n from unittest import mock as bar\n from unittest import mock"
|
||||||
|
location:
|
||||||
|
row: 60
|
||||||
|
column: 4
|
||||||
|
end_location:
|
||||||
|
row: 60
|
||||||
|
column: 41
|
||||||
|
parent: ~
|
||||||
|
- kind: RewriteMockImport
|
||||||
|
location:
|
||||||
|
row: 63
|
||||||
|
column: 11
|
||||||
|
end_location:
|
||||||
|
row: 63
|
||||||
|
column: 22
|
||||||
|
fix:
|
||||||
|
content: "import os\n from unittest import mock as foo\n from unittest import mock as bar\n from unittest import mock"
|
||||||
|
location:
|
||||||
|
row: 63
|
||||||
|
column: 4
|
||||||
|
end_location:
|
||||||
|
row: 63
|
||||||
|
column: 45
|
||||||
|
parent: ~
|
||||||
|
- kind: RewriteMockImport
|
||||||
|
location:
|
||||||
|
row: 63
|
||||||
|
column: 24
|
||||||
|
end_location:
|
||||||
|
row: 63
|
||||||
|
column: 35
|
||||||
|
fix:
|
||||||
|
content: "import os\n from unittest import mock as foo\n from unittest import mock as bar\n from unittest import mock"
|
||||||
|
location:
|
||||||
|
row: 63
|
||||||
|
column: 4
|
||||||
|
end_location:
|
||||||
|
row: 63
|
||||||
|
column: 45
|
||||||
|
parent: ~
|
||||||
|
- kind: RewriteMockImport
|
||||||
|
location:
|
||||||
|
row: 63
|
||||||
|
column: 37
|
||||||
|
end_location:
|
||||||
|
row: 63
|
||||||
|
column: 41
|
||||||
|
fix:
|
||||||
|
content: "import os\n from unittest import mock as foo\n from unittest import mock as bar\n from unittest import mock"
|
||||||
|
location:
|
||||||
|
row: 63
|
||||||
|
column: 4
|
||||||
|
end_location:
|
||||||
|
row: 63
|
||||||
|
column: 45
|
||||||
|
parent: ~
|
||||||
|
- kind: RewriteMockImport
|
||||||
|
location:
|
||||||
|
row: 67
|
||||||
|
column: 4
|
||||||
|
end_location:
|
||||||
|
row: 67
|
||||||
|
column: 51
|
||||||
|
fix:
|
||||||
|
content: "from unittest import mock as foo\n from unittest import mock as bar\n from unittest import mock"
|
||||||
|
location:
|
||||||
|
row: 67
|
||||||
|
column: 4
|
||||||
|
end_location:
|
||||||
|
row: 67
|
||||||
|
column: 51
|
||||||
|
parent: ~
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue