-
Notifications
You must be signed in to change notification settings - Fork 1.1k
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.
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
[pyupgrade
] Replace str,Enum with StrEnum UP042
#10713
Changes from all commits
d4f3396
c2adfce
125d468
4714913
cee7d1b
d677f25
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
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): ... | ||
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`") | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We prefer to omit the fix in the rule message. |
||
} | ||
|
||
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() { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I changed this to a |
||
["", "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(), | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The use of |
||
); | ||
|
||
// 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); | ||
} |
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
(I ran
ruff format
over the fixture. It's not a hard rule, since fixtures are often intentionally unformatted, but it's nice when we don't rely on the formatting anywhere.)