Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
## Summary Add new rule `pyupgrade - UP042` (I picked next available number). Closes #3867 Closes #9569 It should warn + provide a fix `class A(str, Enum)` -> `class A(StrEnum)` for py311+. ## Test Plan Added UP042.py test. ## Notes I did not find a way to call `remove_argument` 2 times consecutively, so the automatic fixing works only for classes that inherit exactly `str, Enum` (regardless of the order). I also plan to extend this rule to support IntEnum in next PR.
- Loading branch information
Showing
8 changed files
with
238 additions
and
0 deletions.
There are no files selected for viewing
13 changes: 13 additions & 0 deletions
13
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,13 @@ | ||
from enum import Enum | ||
|
||
|
||
class A(str, Enum): ... | ||
|
||
|
||
class B(Enum, str): ... | ||
|
||
|
||
class D(int, str, Enum): ... | ||
|
||
|
||
class E(str, int, Enum): ... |
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
160 changes: 160 additions & 0 deletions
160
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,160 @@ | ||
use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation}; | ||
use ruff_macros::{derive_message_formats, violation}; | ||
use ruff_python_ast as ast; | ||
use ruff_python_ast::identifier::Identifier; | ||
use ruff_text_size::Ranged; | ||
|
||
use crate::checkers::ast::Checker; | ||
use crate::importer::ImportRequest; | ||
|
||
/// ## What it does | ||
/// Checks for classes that inherit from both `str` and `enum.Enum`. | ||
/// | ||
/// ## Why is this bad? | ||
/// Python 3.11 introduced `enum.StrEnum`, which is preferred over inheriting | ||
/// from both `str` and `enum.Enum`. | ||
/// | ||
/// ## Example | ||
/// | ||
/// ```python | ||
/// import enum | ||
/// | ||
/// | ||
/// class Foo(str, enum.Enum): | ||
/// ... | ||
/// ``` | ||
/// | ||
/// Use instead: | ||
/// | ||
/// ```python | ||
/// import enum | ||
/// | ||
/// | ||
/// class Foo(enum.StrEnum): | ||
/// ... | ||
/// ``` | ||
/// | ||
/// ## Fix safety | ||
/// | ||
/// Python 3.11 introduced a [breaking change] for enums that inherit from both | ||
/// `str` and `enum.Enum`. Consider the following enum: | ||
/// | ||
/// ```python | ||
/// from enum import Enum | ||
/// | ||
/// | ||
/// class Foo(str, Enum): | ||
/// BAR = "bar" | ||
/// ``` | ||
/// | ||
/// In Python 3.11, the formatted representation of `Foo.BAR` changed as | ||
/// follows: | ||
/// | ||
/// ```python | ||
/// # Python 3.10 | ||
/// f"{Foo.BAR}" # > bar | ||
/// # Python 3.11 | ||
/// f"{Foo.BAR}" # > Foo.BAR | ||
/// ``` | ||
/// | ||
/// Migrating from `str` and `enum.Enum` to `enum.StrEnum` will restore the | ||
/// previous behavior, such that: | ||
/// | ||
/// ```python | ||
/// from enum import StrEnum | ||
/// | ||
/// | ||
/// class Foo(StrEnum): | ||
/// BAR = "bar" | ||
/// | ||
/// | ||
/// f"{Foo.BAR}" # > bar | ||
/// ``` | ||
/// | ||
/// As such, migrating to `enum.StrEnum` will introduce a behavior change for | ||
/// code that relies on the Python 3.11 behavior. | ||
/// | ||
/// ## References | ||
/// - [enum.StrEnum](https://docs.python.org/3/library/enum.html#enum.StrEnum) | ||
/// | ||
/// [breaking change]: https://blog.pecar.me/python-enum | ||
|
||
#[violation] | ||
pub struct ReplaceStrEnum { | ||
name: String, | ||
} | ||
|
||
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`") | ||
} | ||
|
||
fn fix_title(&self) -> Option<String> { | ||
Some("Inherit from `enum.StrEnum`".to_string()) | ||
} | ||
} | ||
|
||
/// 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; | ||
}; | ||
|
||
// Determine whether the class inherits from both `str` and `enum.Enum`. | ||
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) { | ||
match qualified_name.segments() { | ||
["", "str"] => inherits_str = true, | ||
["enum", "Enum"] => inherits_enum = true, | ||
_ => {} | ||
} | ||
} | ||
|
||
// Short-circuit if both `str` and `enum.Enum` are found. | ||
if inherits_str && inherits_enum { | ||
break; | ||
} | ||
} | ||
|
||
// If the class does not inherit both `str` and `enum.Enum`, exit early. | ||
if !inherits_str || !inherits_enum { | ||
return; | ||
}; | ||
|
||
let mut diagnostic = Diagnostic::new( | ||
ReplaceStrEnum { | ||
name: class_def.name.to_string(), | ||
}, | ||
class_def.identifier(), | ||
); | ||
|
||
// If the base classes are _exactly_ `str` and `enum.Enum`, apply a fix. | ||
// TODO(charlie): As an alternative, we could remove both arguments, and replace one of the two | ||
// with `StrEnum`. However, `remove_argument` can't be applied multiple times within a single | ||
// fix; doing so leads to a syntax error. | ||
if arguments.len() == 2 { | ||
diagnostic.try_set_fix(|| { | ||
let (import_edit, binding) = checker.importer().get_or_import_symbol( | ||
&ImportRequest::import("enum", "StrEnum"), | ||
class_def.start(), | ||
checker.semantic(), | ||
)?; | ||
Ok(Fix::unsafe_edits( | ||
import_edit, | ||
[Edit::range_replacement( | ||
format!("({binding})"), | ||
arguments.range(), | ||
)], | ||
)) | ||
}); | ||
} | ||
|
||
checker.diagnostics.push(diagnostic); | ||
} |
55 changes: 55 additions & 0 deletions
55
..._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,55 @@ | ||
--- | ||
source: crates/ruff_linter/src/rules/pyupgrade/mod.rs | ||
--- | ||
UP042.py:4:7: UP042 [*] Class A inherits from both `str` and `enum.Enum` | ||
| | ||
4 | class A(str, Enum): ... | ||
| ^ UP042 | ||
| | ||
= help: Inherit from `enum.StrEnum` | ||
|
||
ℹ Unsafe fix | ||
1 |-from enum import Enum | ||
1 |+from enum import Enum, StrEnum | ||
2 2 | | ||
3 3 | | ||
4 |-class A(str, Enum): ... | ||
4 |+class A(StrEnum): ... | ||
5 5 | | ||
6 6 | | ||
7 7 | class B(Enum, str): ... | ||
|
||
UP042.py:7:7: UP042 [*] Class B inherits from both `str` and `enum.Enum` | ||
| | ||
7 | class B(Enum, str): ... | ||
| ^ UP042 | ||
| | ||
= help: Inherit from `enum.StrEnum` | ||
|
||
ℹ Unsafe fix | ||
1 |-from enum import Enum | ||
1 |+from enum import Enum, StrEnum | ||
2 2 | | ||
3 3 | | ||
4 4 | class A(str, Enum): ... | ||
5 5 | | ||
6 6 | | ||
7 |-class B(Enum, str): ... | ||
7 |+class B(StrEnum): ... | ||
8 8 | | ||
9 9 | | ||
10 10 | class D(int, str, Enum): ... | ||
|
||
UP042.py:10:7: UP042 Class D inherits from both `str` and `enum.Enum` | ||
| | ||
10 | class D(int, str, Enum): ... | ||
| ^ UP042 | ||
| | ||
= help: Inherit from `enum.StrEnum` | ||
|
||
UP042.py:13:7: UP042 Class E inherits from both `str` and `enum.Enum` | ||
| | ||
13 | class E(str, int, Enum): ... | ||
| ^ UP042 | ||
| | ||
= help: Inherit from `enum.StrEnum` |
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Oops, something went wrong.