mirror of
https://github.com/astral-sh/ruff.git
synced 2025-09-29 13:25:17 +00:00
[pyupgrade
] Replace str, Enum
with StrEnum
(UP042
) (#10713)
## Summary Add new rule `pyupgrade - UP042` (I picked next available number). Closes https://github.com/astral-sh/ruff/discussions/3867 Closes https://github.com/astral-sh/ruff/issues/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.
This commit is contained in:
parent
323264dec2
commit
b45fd61ec5
8 changed files with 238 additions and 0 deletions
13
crates/ruff_linter/resources/test/fixtures/pyupgrade/UP042.py
vendored
Normal file
13
crates/ruff_linter/resources/test/fixtures/pyupgrade/UP042.py
vendored
Normal file
|
@ -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): ...
|
|
@ -406,6 +406,11 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
|
|||
if checker.enabled(Rule::UselessObjectInheritance) {
|
||||
pyupgrade::rules::useless_object_inheritance(checker, class_def);
|
||||
}
|
||||
if checker.enabled(Rule::ReplaceStrEnum) {
|
||||
if checker.settings.target_version >= PythonVersion::Py311 {
|
||||
pyupgrade::rules::replace_str_enum(checker, class_def);
|
||||
}
|
||||
}
|
||||
if checker.enabled(Rule::UnnecessaryClassParentheses) {
|
||||
pyupgrade::rules::unnecessary_class_parentheses(checker, class_def);
|
||||
}
|
||||
|
|
|
@ -547,6 +547,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
|
|||
(Pyupgrade, "039") => (RuleGroup::Stable, rules::pyupgrade::rules::UnnecessaryClassParentheses),
|
||||
(Pyupgrade, "040") => (RuleGroup::Stable, rules::pyupgrade::rules::NonPEP695TypeAlias),
|
||||
(Pyupgrade, "041") => (RuleGroup::Stable, rules::pyupgrade::rules::TimeoutErrorAlias),
|
||||
(Pyupgrade, "042") => (RuleGroup::Preview, rules::pyupgrade::rules::ReplaceStrEnum),
|
||||
|
||||
// pydocstyle
|
||||
(Pydocstyle, "100") => (RuleGroup::Stable, rules::pydocstyle::rules::UndocumentedPublicModule),
|
||||
|
|
|
@ -61,6 +61,7 @@ mod tests {
|
|||
#[test_case(Rule::ReplaceUniversalNewlines, Path::new("UP021.py"))]
|
||||
#[test_case(Rule::SuperCallWithParameters, Path::new("UP008.py"))]
|
||||
#[test_case(Rule::TimeoutErrorAlias, Path::new("UP041.py"))]
|
||||
#[test_case(Rule::ReplaceStrEnum, Path::new("UP042.py"))]
|
||||
#[test_case(Rule::TypeOfPrimitive, Path::new("UP003.py"))]
|
||||
#[test_case(Rule::TypingTextStrAlias, Path::new("UP019.py"))]
|
||||
#[test_case(Rule::UTF8EncodingDeclaration, Path::new("UP009_0.py"))]
|
||||
|
|
|
@ -18,6 +18,7 @@ pub(crate) use printf_string_formatting::*;
|
|||
pub(crate) use quoted_annotation::*;
|
||||
pub(crate) use redundant_open_modes::*;
|
||||
pub(crate) use replace_stdout_stderr::*;
|
||||
pub(crate) use replace_str_enum::*;
|
||||
pub(crate) use replace_universal_newlines::*;
|
||||
pub(crate) use super_call_with_parameters::*;
|
||||
pub(crate) use timeout_error_alias::*;
|
||||
|
@ -58,6 +59,7 @@ mod printf_string_formatting;
|
|||
mod quoted_annotation;
|
||||
mod redundant_open_modes;
|
||||
mod replace_stdout_stderr;
|
||||
mod replace_str_enum;
|
||||
mod replace_universal_newlines;
|
||||
mod super_call_with_parameters;
|
||||
mod timeout_error_alias;
|
||||
|
|
160
crates/ruff_linter/src/rules/pyupgrade/rules/replace_str_enum.rs
Normal file
160
crates/ruff_linter/src/rules/pyupgrade/rules/replace_str_enum.rs
Normal file
|
@ -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);
|
||||
}
|
|
@ -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`
|
1
ruff.schema.json
generated
1
ruff.schema.json
generated
|
@ -3867,6 +3867,7 @@
|
|||
"UP04",
|
||||
"UP040",
|
||||
"UP041",
|
||||
"UP042",
|
||||
"W",
|
||||
"W1",
|
||||
"W19",
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue