[pyupgrade] Replace str,Enum with StrEnum UP042

crates/ruff_linter/resources/test/fixtures/pyupgrade/UP042.py

@@ -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): ...

crates/ruff_linter/src/checkers/ast/analyze/statement.rs

@@ -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);
             }

crates/ruff_linter/src/codes.rs

@@ -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),

crates/ruff_linter/src/rules/pyupgrade/mod.rs

@@ -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"))]

crates/ruff_linter/src/rules/pyupgrade/rules/mod.rs

@@ -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;

crates/ruff_linter/src/rules/pyupgrade/rules/replace_str_enum.rs

@@ -0,0 +1,155 @@
+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);
+}

crates/ruff_linter/src/rules/pyupgrade/snapshots/ruff_linter__rules__pyupgrade__tests__UP042.py.snap

@@ -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`