Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Rewrite mock import with starred imports #3566

Merged
merged 1 commit into from
Mar 17, 2023
Merged

Rewrite mock import with starred imports #3566

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
31 changes: 18 additions & 13 deletions crates/ruff/resources/test/fixtures/pyupgrade/UP026.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,24 +1,29 @@
# These should be changed
# Error (`from unittest import mock`)
if True:
import mock

# Error (`from unittest import mock`)
if True:
import mock, sys

# This goes to from unitest import mock
# Error (`from unittest.mock import *`)
if True:
from mock import *

# Error (`from unittest import mock`)
import mock.mock

# Mock should go on a new line as `from unittest import mock`
# Error (`from unittest import mock`)
import contextlib, mock, sys

# Mock should go on a new line as `from unittest import mock`
# Error (`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`
# Error (`from unittest import mock`)
from mock import mock

# Should keep trailing comma
# Error (keep trailing comma)
from mock import (
mock,
a,
Expand All @@ -32,7 +37,7 @@
mock,
)

# Should not get a trailing comma
# Error (avoid trailing comma)
from mock import (
mock,
a,
Expand All @@ -57,16 +62,16 @@
c
)

# These should not change:
# OK
import os, io

# Mock should go on a new line as `from unittest import mock`
# Error (`from unittest import mock`)
import mock, mock

# Mock should go on a new line as `from unittest import mock as foo`
# Error (`from unittest import mock as foo`)
import mock as foo

# Mock should go on a new line as `from unittest import mock as foo`
# Error (`from unittest import mock as foo`)
from mock import mock as foo

if True:
Expand All @@ -81,8 +86,8 @@
from mock import mock as foo, mock as bar, mock


# This should be unchanged.
# OK.
x = mock.Mock()

# This should change to `mock.Mock`().
# Error (`mock.Mock()`).
x = mock.mock.Mock()
91 changes: 53 additions & 38 deletions crates/ruff/src/rules/pyupgrade/rules/rewrite_mock_import.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,22 +93,6 @@ fn clean_import_aliases(aliases: Vec<ImportAlias>) -> (Vec<ImportAlias>, Vec<Opt
(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: &Stylist) -> String {
let mut content = String::new();
for alias in aliases {
Expand Down Expand Up @@ -180,20 +164,12 @@ fn format_import_from(
let mut tree = match_module(module_text).unwrap();
let mut import = match_import_from(&mut tree)?;

let ImportFrom {
names: ImportNames::Aliases(aliases),
if let ImportFrom {
names: ImportNames::Star(..),
..
} = import.clone() else {
unreachable!("Expected ImportNames::Aliases");
};

let has_mock_member = includes_mock_member(&aliases);
let (clean_aliases, mock_aliases) = clean_import_aliases(aliases);

Ok(if clean_aliases.is_empty() {
format_mocks(mock_aliases, indent, stylist)
} else {
import.names = ImportNames::Aliases(clean_aliases);
} = import
{
// Ex) `from mock import *`
import.module = Some(NameOrAttribute::A(Box::new(Attribute {
value: Box::new(Expression::Name(Box::new(Name {
value: "unittest",
Expand All @@ -212,22 +188,61 @@ fn format_import_from(
lpar: vec![],
rpar: vec![],
})));

let mut state = CodegenState {
default_newline: stylist.line_ending(),
default_indent: stylist.indentation(),
..CodegenState::default()
};
tree.codegen(&mut state);
Ok(state.to_string())
} else if let ImportFrom {
names: ImportNames::Aliases(aliases),
..
} = import
{
// Ex) `from mock import mock`
let (clean_aliases, mock_aliases) = clean_import_aliases(aliases.clone());
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 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
})
let mut state = CodegenState {
default_newline: stylist.line_ending(),
default_indent: stylist.indentation(),
..CodegenState::default()
};
tree.codegen(&mut state);

let mut content = state.to_string();
if !mock_aliases.is_empty() {
content.push_str(stylist.line_ending());
content.push_str(indent);
content.push_str(&format_mocks(mock_aliases, indent, stylist));
}
content
})
} else {
unreachable!("Expected ImportNames::Aliases | ImportNames::Star");
}
}

/// UP026
Expand Down