-
Notifications
You must be signed in to change notification settings - Fork 892
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[
pyupgrade
] Replace str,Enum with StrEnum UP042
Basically it closes #3867
- Loading branch information
Showing
8 changed files
with
229 additions
and
0 deletions.
There are no files selected for viewing
15 changes: 15 additions & 0 deletions
15
crates/ruff_linter/resources/test/fixtures/pyupgrade/UP042.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,15 @@ | ||
from enum import Enum | ||
|
||
class A(str, Enum): | ||
... | ||
|
||
class B(Enum, str): | ||
... | ||
|
||
class D(int, str, Enum): | ||
... | ||
|
||
class E(str, int, Enum): | ||
... | ||
|
||
# TODO: add more cases |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
123 changes: 123 additions & 0 deletions
123
crates/ruff_linter/src/rules/pyupgrade/rules/replace_str_enum.rs
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,123 @@ | ||
use crate::checkers::ast::Checker; | ||
use crate::importer::ImportRequest; | ||
use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation}; | ||
use ruff_macros::{derive_message_formats, violation}; | ||
use ruff_python_ast::{self as ast}; | ||
use ruff_text_size::Ranged; | ||
|
||
/// ## What it does | ||
/// Checks for classes that inherit from both `str` and `enum.Enum`. | ||
/// | ||
/// ## Why is this bad? | ||
/// Since Python 3.11, `enum.StrEnum` exists and is preferred over | ||
/// inheriting from `str` and `enum.Enum`. | ||
/// | ||
/// ## Example | ||
/// ```python | ||
/// class Foo(str, enum.Enum): | ||
/// ... | ||
/// ``` | ||
/// | ||
/// Use instead: | ||
/// ```python | ||
/// class Foo(enum.StrEnum): | ||
/// ... | ||
/// ``` | ||
/// | ||
/// ## References | ||
/// - [enum.StrEnum](https://docs.python.org/3/library/enum.html#enum.StrEnum) | ||
|
||
#[violation] | ||
pub struct ReplaceStrEnum { | ||
name: String, | ||
args_len: usize, | ||
} | ||
|
||
impl Violation for ReplaceStrEnum { | ||
const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes; | ||
|
||
#[derive_message_formats] | ||
fn message(&self) -> String { | ||
let ReplaceStrEnum { name, .. } = self; | ||
format!( | ||
"Class {name} inherits from both `str` and `enum.Enum`. Prefer `enum.StrEnum` instead." | ||
) | ||
} | ||
|
||
fn fix_title(&self) -> Option<String> { | ||
let ReplaceStrEnum { args_len, .. } = self; | ||
|
||
if *args_len == 2 { | ||
Some("Replace `str` and `enum.Enum` with `enum.StrEnum`".to_string()) | ||
} else { | ||
None | ||
} | ||
} | ||
} | ||
|
||
/// UP042 | ||
pub(crate) fn replace_str_enum(checker: &mut Checker, class_def: &ast::StmtClassDef) { | ||
let Some(arguments) = class_def.arguments.as_deref() else { | ||
// class does not inherit anything, exit early | ||
return; | ||
}; | ||
|
||
let mut inherits_str = false; | ||
let mut inherits_enum = false; | ||
for base in arguments.args.iter() { | ||
if let Some(qualified_name) = checker.semantic().resolve_qualified_name(base) { | ||
if matches!(qualified_name.segments(), ["", "str"]) { | ||
inherits_str = true; | ||
} else if matches!(qualified_name.segments(), ["enum", "Enum"]) { | ||
inherits_enum = true; | ||
} | ||
} | ||
|
||
if inherits_str && inherits_enum { | ||
// no need to check other inherited classes, we found both str & enum.Enum | ||
break; | ||
} | ||
} | ||
|
||
if !inherits_str || !inherits_enum { | ||
// exit early if class does not inherit both str and enum.Enum | ||
return; | ||
}; | ||
|
||
let mut diagnostic = Diagnostic::new( | ||
ReplaceStrEnum { | ||
name: class_def.name.to_string(), | ||
args_len: arguments.len(), | ||
}, | ||
class_def.range(), | ||
); | ||
|
||
if arguments.len() == 2 { | ||
// a fix is available only for classes that inherit exactly 2 arguments: str, Enum, | ||
// because `remove_argument` cannot be called multiple times consecutively... | ||
// for classes that inherit str, Enum and something else, generate a warning. | ||
diagnostic.try_set_fix(|| { | ||
let (import_edit, binding) = checker.importer().get_or_import_symbol( | ||
&ImportRequest::import("enum", "StrEnum"), | ||
class_def.start(), | ||
checker.semantic(), | ||
)?; | ||
|
||
// `binding` here is `StrEnum`. | ||
|
||
// class inherits exactly 2 arguments. | ||
// replace all `(str, Enum)` arguments with `(StrEnum)`. | ||
let fix = Fix::unsafe_edits( | ||
import_edit, | ||
[Edit::range_replacement( | ||
format!("({binding})"), | ||
arguments.range(), | ||
)], | ||
); | ||
|
||
Ok(fix) | ||
}); | ||
} | ||
|
||
checker.diagnostics.push(diagnostic); | ||
} |
71 changes: 71 additions & 0 deletions
71
..._linter/src/rules/pyupgrade/snapshots/ruff_linter__rules__pyupgrade__tests__UP042.py.snap
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,71 @@ | ||
--- | ||
source: crates/ruff_linter/src/rules/pyupgrade/mod.rs | ||
--- | ||
UP042.py:3:1: UP042 [*] Class A inherits from both `str` and `enum.Enum`. Prefer `enum.StrEnum` instead. | ||
| | ||
1 | from enum import Enum | ||
2 | | ||
3 | / class A(str, Enum): | ||
4 | | ... | ||
| |_______^ UP042 | ||
5 | | ||
6 | class B(Enum, str): | ||
| | ||
= help: Replace `str` and `enum.Enum` with `enum.StrEnum` | ||
|
||
ℹ Unsafe fix | ||
1 |-from enum import Enum | ||
1 |+from enum import Enum, StrEnum | ||
2 2 | | ||
3 |-class A(str, Enum): | ||
3 |+class A(StrEnum): | ||
4 4 | ... | ||
5 5 | | ||
6 6 | class B(Enum, str): | ||
|
||
UP042.py:6:1: UP042 [*] Class B inherits from both `str` and `enum.Enum`. Prefer `enum.StrEnum` instead. | ||
| | ||
4 | ... | ||
5 | | ||
6 | / class B(Enum, str): | ||
7 | | ... | ||
| |_______^ UP042 | ||
8 | | ||
9 | class D(int, str, Enum): | ||
| | ||
= help: Replace `str` and `enum.Enum` with `enum.StrEnum` | ||
|
||
ℹ Unsafe fix | ||
1 |-from enum import Enum | ||
1 |+from enum import Enum, StrEnum | ||
2 2 | | ||
3 3 | class A(str, Enum): | ||
4 4 | ... | ||
5 5 | | ||
6 |-class B(Enum, str): | ||
6 |+class B(StrEnum): | ||
7 7 | ... | ||
8 8 | | ||
9 9 | class D(int, str, Enum): | ||
|
||
UP042.py:9:1: UP042 Class D inherits from both `str` and `enum.Enum`. Prefer `enum.StrEnum` instead. | ||
| | ||
7 | ... | ||
8 | | ||
9 | / class D(int, str, Enum): | ||
10 | | ... | ||
| |_______^ UP042 | ||
11 | | ||
12 | class E(str, int, Enum): | ||
| | ||
|
||
UP042.py:12:1: UP042 Class E inherits from both `str` and `enum.Enum`. Prefer `enum.StrEnum` instead. | ||
| | ||
10 | ... | ||
11 | | ||
12 | / class E(str, int, Enum): | ||
13 | | ... | ||
| |_______^ UP042 | ||
14 | | ||
15 | # TODO: add more cases | ||
| |