[pyupgrade] Split UP007 to two individual rules for Union and `…

…Optional` (`UP007`, `UP045`) (#15313)

Co-authored-by: Micha Reiser <[email protected]>

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

@@ -1,16 +1,7 @@
 import typing
-from typing import Optional
 from typing import Union
 
 
-def f(x: Optional[str]) -> None:
-    ...
-
-
-def f(x: typing.Optional[str]) -> None:
-    ...
-
-
 def f(x: Union[str, int, Union[float, bytes]]) -> None:
     ...
 
@@ -52,9 +43,6 @@ def f(x: Union[("str", "int"), float]) -> None:
 
 
 def f() -> None:
-    x: Optional[str]
-    x = Optional[str]
-
     x = Union[str, int]
     x = Union["str", "int"]
     x: Union[str, int]
@@ -85,31 +73,6 @@ def f(x: Union[str, lambda: int]) -> None:
     ...
 
 
-def f(x: Optional[int : float]) -> None:
-    ...
-
-
-def f(x: Optional[str, int : float]) -> None:
-    ...
-
-
-def f(x: Optional[int, float]) -> None:
-    ...
-
-
-# Regression test for: https://github.com/astral-sh/ruff/issues/7131
-class ServiceRefOrValue:
-    service_specification: Optional[
-        list[ServiceSpecificationRef]
-        | list[ServiceSpecification]
-    ] = None
-
-
-# Regression test for: https://github.com/astral-sh/ruff/issues/7201
-class ServiceRefOrValue:
-    service_specification: Optional[str]is not True = None
-
-
 # Regression test for: https://github.com/astral-sh/ruff/issues/7452
 class Collection(Protocol[*_B0]):
     def __iter__(self) -> Iterator[Union[*_B0]]:

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

@@ -0,0 +1,44 @@
+import typing
+from typing import Optional
+
+
+def f(x: Optional[str]) -> None:
+    ...
+
+
+def f(x: typing.Optional[str]) -> None:
+    ...
+
+
+def f() -> None:
+    x: Optional[str]
+    x = Optional[str]
+
+
+def f(x: list[Optional[int]]) -> None:
+    ...
+
+
+def f(x: Optional[int : float]) -> None:
+    ...
+
+
+def f(x: Optional[str, int : float]) -> None:
+    ...
+
+
+def f(x: Optional[int, float]) -> None:
+    ...
+
+
+# Regression test for: https://github.com/astral-sh/ruff/issues/7131
+class ServiceRefOrValue:
+    service_specification: Optional[
+        list[ServiceSpecificationRef]
+        | list[ServiceSpecification]
+    ] = None
+
+
+# Regression test for: https://github.com/astral-sh/ruff/issues/7201
+class ServiceRefOrValue:
+    service_specification: Optional[str]is not True = None

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

@@ -27,7 +27,8 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
             // Ex) Optional[...], Union[...]
             if checker.any_enabled(&[
                 Rule::FutureRewritableTypeAnnotation,
-                Rule::NonPEP604Annotation,
+                Rule::NonPEP604AnnotationUnion,
+                Rule::NonPEP604AnnotationOptional,
             ]) {
                 if let Some(operator) = typing::to_pep604_operator(value, slice, &checker.semantic)
                 {
@@ -43,15 +44,18 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                             );
                         }
                     }
-                    if checker.enabled(Rule::NonPEP604Annotation) {
+                    if checker.any_enabled(&[
+                        Rule::NonPEP604AnnotationUnion,
+                        Rule::NonPEP604AnnotationOptional,
+                    ]) {
                         if checker.source_type.is_stub()
                             || checker.settings.target_version >= PythonVersion::Py310
                             || (checker.settings.target_version >= PythonVersion::Py37
                                 && checker.semantic.future_annotations_or_stub()
                                 && checker.semantic.in_annotation()
                                 && !checker.settings.pyupgrade.keep_runtime_typing)
                         {
-                            pyupgrade::rules::use_pep604_annotation(checker, expr, slice, operator);
+                            pyupgrade::rules::non_pep604_annotation(checker, expr, slice, operator);
                         }
                     }
                 }

crates/ruff_linter/src/codes.rs

@@ -501,7 +501,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
         (Pyupgrade, "004") => (RuleGroup::Stable, rules::pyupgrade::rules::UselessObjectInheritance),
         (Pyupgrade, "005") => (RuleGroup::Stable, rules::pyupgrade::rules::DeprecatedUnittestAlias),
         (Pyupgrade, "006") => (RuleGroup::Stable, rules::pyupgrade::rules::NonPEP585Annotation),
-        (Pyupgrade, "007") => (RuleGroup::Stable, rules::pyupgrade::rules::NonPEP604Annotation),
+        (Pyupgrade, "007") => (RuleGroup::Stable, rules::pyupgrade::rules::NonPEP604AnnotationUnion),
         (Pyupgrade, "008") => (RuleGroup::Stable, rules::pyupgrade::rules::SuperCallWithParameters),
         (Pyupgrade, "009") => (RuleGroup::Stable, rules::pyupgrade::rules::UTF8EncodingDeclaration),
         (Pyupgrade, "010") => (RuleGroup::Stable, rules::pyupgrade::rules::UnnecessaryFutureImport),
@@ -538,6 +538,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
         (Pyupgrade, "042") => (RuleGroup::Preview, rules::pyupgrade::rules::ReplaceStrEnum),
         (Pyupgrade, "043") => (RuleGroup::Stable, rules::pyupgrade::rules::UnnecessaryDefaultTypeArgs),
         (Pyupgrade, "044") => (RuleGroup::Preview, rules::pyupgrade::rules::NonPEP646Unpack),
+        (Pyupgrade, "045") => (RuleGroup::Preview, rules::pyupgrade::rules::NonPEP604AnnotationOptional),
 
         // pydocstyle
         (Pydocstyle, "100") => (RuleGroup::Stable, rules::pydocstyle::rules::UndocumentedPublicModule),

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

@@ -39,7 +39,8 @@ mod tests {
     #[test_case(Rule::NonPEP585Annotation, Path::new("UP006_1.py"))]
     #[test_case(Rule::NonPEP585Annotation, Path::new("UP006_2.py"))]
     #[test_case(Rule::NonPEP585Annotation, Path::new("UP006_3.py"))]
-    #[test_case(Rule::NonPEP604Annotation, Path::new("UP007.py"))]
+    #[test_case(Rule::NonPEP604AnnotationUnion, Path::new("UP007.py"))]
+    #[test_case(Rule::NonPEP604AnnotationUnion, Path::new("UP045.py"))]
     #[test_case(Rule::NonPEP604Isinstance, Path::new("UP038.py"))]
     #[test_case(Rule::OSErrorAlias, Path::new("UP024_0.py"))]
     #[test_case(Rule::OSErrorAlias, Path::new("UP024_1.py"))]
@@ -111,6 +112,19 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn up007_preview() -> Result<()> {
+        let diagnostics = test_path(
+            Path::new("pyupgrade/UP045.py"),
+            &settings::LinterSettings {
+                preview: PreviewMode::Enabled,
+                ..settings::LinterSettings::for_rule(Rule::NonPEP604AnnotationUnion)
+            },
+        )?;
+        assert_messages!(diagnostics);
+        Ok(())
+    }
+
     #[test]
     fn async_timeout_error_alias_not_applied_py310() -> Result<()> {
         let diagnostics = test_path(
@@ -201,7 +215,10 @@ mod tests {
             Path::new("pyupgrade/future_annotations.py"),
             &settings::LinterSettings {
                 target_version: PythonVersion::Py37,
-                ..settings::LinterSettings::for_rule(Rule::NonPEP604Annotation)
+                ..settings::LinterSettings::for_rules([
+                    Rule::NonPEP604AnnotationUnion,
+                    Rule::NonPEP604AnnotationOptional,
+                ])
             },
         )?;
         assert_messages!(diagnostics);
@@ -214,7 +231,10 @@ mod tests {
             Path::new("pyupgrade/future_annotations.py"),
             &settings::LinterSettings {
                 target_version: PythonVersion::Py310,
-                ..settings::LinterSettings::for_rule(Rule::NonPEP604Annotation)
+                ..settings::LinterSettings::for_rules([
+                    Rule::NonPEP604AnnotationUnion,
+                    Rule::NonPEP604AnnotationOptional,
+                ])
             },
         )?;
         assert_messages!(diagnostics);

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

@@ -1,13 +1,16 @@
-use ruff_diagnostics::{Applicability, Diagnostic, Edit, Fix, FixAvailability, Violation};
+use ruff_diagnostics::{
+    Applicability, Diagnostic, DiagnosticKind, Edit, Fix, FixAvailability, Violation,
+};
 use ruff_macros::{derive_message_formats, ViolationMetadata};
 use ruff_python_ast::helpers::{pep_604_optional, pep_604_union};
 use ruff_python_ast::{self as ast, Expr};
 use ruff_python_semantic::analyze::typing::Pep604Operator;
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
+use crate::codes::Rule;
 use crate::fix::edits::pad;
-use crate::settings::types::PythonVersion;
+use crate::settings::types::{PreviewMode, PythonVersion};
 
 /// ## What it does
 /// Check for type annotations that can be rewritten based on [PEP 604] syntax.
@@ -37,6 +40,10 @@ use crate::settings::types::PythonVersion;
 /// foo: int | str = 1
 /// ```
 ///
+/// ## Preview
+/// In preview mode, this rule only checks for usages of `typing.Union`,
+/// while `UP045` checks for `typing.Optional`.
+///
 /// ## Fix safety
 /// This rule's fix is marked as unsafe, as it may lead to runtime errors when
 /// alongside libraries that rely on runtime type annotations, like Pydantic,
@@ -50,9 +57,9 @@ use crate::settings::types::PythonVersion;
 ///
 /// [PEP 604]: https://peps.python.org/pep-0604/
 #[derive(ViolationMetadata)]
-pub(crate) struct NonPEP604Annotation;
+pub(crate) struct NonPEP604AnnotationUnion;
 
-impl Violation for NonPEP604Annotation {
+impl Violation for NonPEP604AnnotationUnion {
     const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes;
 
     #[derive_message_formats]
@@ -65,8 +72,64 @@ impl Violation for NonPEP604Annotation {
     }
 }
 
-/// UP007
-pub(crate) fn use_pep604_annotation(
+/// ## What it does
+/// Check for `typing.Optional` annotations that can be rewritten based on [PEP 604] syntax.
+///
+/// ## Why is this bad?
+/// [PEP 604] introduced a new syntax for union type annotations based on the
+/// `|` operator. This syntax is more concise and readable than the previous
+/// `typing.Optional` syntax.
+///
+/// This rule is enabled when targeting Python 3.10 or later (see:
+/// [`target-version`]). By default, it's _also_ enabled for earlier Python
+/// versions if `from __future__ import annotations` is present, as
+/// `__future__` annotations are not evaluated at runtime. If your code relies
+/// on runtime type annotations (either directly or via a library like
+/// Pydantic), you can disable this behavior for Python versions prior to 3.10
+/// by setting [`lint.pyupgrade.keep-runtime-typing`] to `true`.
+///
+/// ## Example
+/// ```python
+/// from typing import Optional
+///
+/// foo: Optional[int] = None
+/// ```
+///
+/// Use instead:
+/// ```python
+/// foo: int | None = None
+/// ```
+///
+/// ## Fix safety
+/// This rule's fix is marked as unsafe, as it may lead to runtime errors when
+/// alongside libraries that rely on runtime type annotations, like Pydantic,
+/// on Python versions prior to Python 3.10. It may also lead to runtime errors
+/// in unusual and likely incorrect type annotations where the type does not
+/// support the `|` operator.
+///
+/// ## Options
+/// - `target-version`
+/// - `lint.pyupgrade.keep-runtime-typing`
+///
+/// [PEP 604]: https://peps.python.org/pep-0604/
+#[derive(ViolationMetadata)]
+pub(crate) struct NonPEP604AnnotationOptional;
+
+impl Violation for NonPEP604AnnotationOptional {
+    const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes;
+
+    #[derive_message_formats]
+    fn message(&self) -> String {
+        "Use `X | None` for type annotations".to_string()
+    }
+
+    fn fix_title(&self) -> Option<String> {
+        Some("Convert to `X | None`".to_string())
+    }
+}
+
+/// UP007, UP045
+pub(crate) fn non_pep604_annotation(
     checker: &mut Checker,
     expr: &Expr,
     slice: &Expr,
@@ -86,7 +149,23 @@ pub(crate) fn use_pep604_annotation(
 
     match operator {
         Pep604Operator::Optional => {
-            let mut diagnostic = Diagnostic::new(NonPEP604Annotation, expr.range());
+            let (rule, diagnostic_kind) = match checker.settings.preview {
+                PreviewMode::Disabled => (
+                    Rule::NonPEP604AnnotationUnion,
+                    DiagnosticKind::from(NonPEP604AnnotationUnion),
+                ),
+                PreviewMode::Enabled => (
+                    Rule::NonPEP604AnnotationOptional,
+                    DiagnosticKind::from(NonPEP604AnnotationOptional),
+                ),
+            };
+
+            if !checker.enabled(rule) {
+                return;
+            }
+
+            let mut diagnostic = Diagnostic::new(diagnostic_kind, expr.range());
+
             if fixable {
                 match slice {
                     Expr::Tuple(_) => {
@@ -110,7 +189,11 @@ pub(crate) fn use_pep604_annotation(
             checker.diagnostics.push(diagnostic);
         }
         Pep604Operator::Union => {
-            let mut diagnostic = Diagnostic::new(NonPEP604Annotation, expr.range());
+            if !checker.enabled(Rule::NonPEP604AnnotationUnion) {
+                return;
+            }
+
+            let mut diagnostic = Diagnostic::new(NonPEP604AnnotationUnion, expr.range());
             if fixable {
                 match slice {
                     Expr::Slice(_) => {