Implement PYI030: Unnecessary literal union (#5570)

Implements PYI030 as part of
#848

> Union expressions should never have more than one Literal member, as
Literal[1] | Literal[2] is semantically identical to Literal[1, 2].

Note we differ slightly from the flake8-pyi implementation:

- We detect cases where there are parentheses or nested unions
- We detect cases with mixed `Union` and `|` syntax
- We use the same error message for all violations; flake8-pyi has two
different messages
- We retain the user's quoting style when displaying string literals;
flake8-pyi uses single quotes
- We warn on duplicates of the same literal `Literal[1] | Literal[1]`

crates/ruff/resources/test/fixtures/flake8_pyi/PYI030.py

@@ -0,0 +1,24 @@
+from typing import Literal
+# Shouldn't emit for any cases in the non-stub file for compatibility with flake8-pyi.
+# Note that this rule could be applied here in the future.
+
+field1: Literal[1]  # OK
+field2: Literal[1] | Literal[2]  # OK
+
+def func1(arg1: Literal[1] | Literal[2]):  # OK
+    print(arg1)
+
+
+def func2() -> Literal[1] | Literal[2]:  # OK
+    return "my Literal[1]ing"
+
+
+field3: Literal[1] | Literal[2] | str  # OK
+field4: str | Literal[1] | Literal[2]  # OK
+field5: Literal[1] | str | Literal[2]  # OK
+field6: Literal[1] | bool | Literal[2] | str  # OK
+field7 = Literal[1] | Literal[2]  # OK
+field8: Literal[1] | (Literal[2] | str)  # OK
+field9: Literal[1] | (Literal[2] | str)  # OK
+field10: (Literal[1] | str) | Literal[2]  # OK
+field11: dict[Literal[1] | Literal[2], str]  # OK

crates/ruff/resources/test/fixtures/flake8_pyi/PYI030.pyi

@@ -0,0 +1,86 @@
+import typing
+import typing_extensions
+from typing import Literal
+
+# Shouldn't affect non-union field types.
+field1: Literal[1]  # OK
+
+# Should emit for duplicate field types.
+field2: Literal[1] | Literal[2]  # Error
+
+# Should emit for union types in arguments.
+def func1(arg1: Literal[1] | Literal[2]):  # Error
+    print(arg1)
+
+
+# Should emit for unions in return types.
+def func2() -> Literal[1] | Literal[2]:  # Error
+    return "my Literal[1]ing"
+
+
+# Should emit in longer unions, even if not directly adjacent.
+field3: Literal[1] | Literal[2] | str  # Error
+field4: str | Literal[1] | Literal[2]  # Error
+field5: Literal[1] | str | Literal[2]  # Error
+field6: Literal[1] | bool | Literal[2] | str  # Error
+
+# Should emit for non-type unions.
+field7 = Literal[1] | Literal[2]  # Error
+
+# Should emit for parenthesized unions.
+field8: Literal[1] | (Literal[2] | str)  # Error
+
+# Should handle user parentheses when fixing.
+field9: Literal[1] | (Literal[2] | str)  # Error
+field10: (Literal[1] | str) | Literal[2]  # Error
+
+# Should emit for union in generic parent type.
+field11: dict[Literal[1] | Literal[2], str]  # Error
+
+# Should emit for unions with more than two cases
+field12: Literal[1] | Literal[2] | Literal[3]  # Error
+field13: Literal[1] | Literal[2] | Literal[3] | Literal[4]  # Error
+
+# Should emit for unions with more than two cases, even if not directly adjacent
+field14: Literal[1] | Literal[2] | str | Literal[3]  # Error
+
+# Should emit for unions with mixed literal internal types
+field15: Literal[1] | Literal["foo"] | Literal[True]  # Error
+
+# Shouldn't emit for duplicate field types with same value; covered by Y016
+field16: Literal[1] | Literal[1]  # OK
+
+# Shouldn't emit if in new parent type
+field17: Literal[1] | dict[Literal[2], str]  # OK
+
+# Shouldn't emit if not in a union parent
+field18: dict[Literal[1], Literal[2]]  # OK
+
+# Should respect name of literal type used
+field19: typing.Literal[1] | typing.Literal[2]  # Error
+
+# Should emit in cases with newlines
+field20: typing.Union[
+    Literal[
+        1  # test
+    ],
+    Literal[2],
+]  # Error, newline and comment will not be emitted in message
+
+# Should handle multiple unions with multiple members
+field21: Literal[1, 2] | Literal[3, 4]  # Error
+
+# Should emit in cases with `typing.Union` instead of `|`
+field22: typing.Union[Literal[1], Literal[2]]  # Error
+
+# Should emit in cases with `typing_extensions.Literal`
+field23: typing_extensions.Literal[1] | typing_extensions.Literal[2]  # Error
+
+# Should emit in cases with nested `typing.Union`
+field24: typing.Union[Literal[1], typing.Union[Literal[2], str]]  # Error
+
+# Should emit in cases with mixed `typing.Union` and `|`
+field25: typing.Union[Literal[1], Literal[2] | str]  # Error
+
+# Should emit only once in cases with multiple nested `typing.Union`
+field24: typing.Union[Literal[1], typing.Union[Literal[2], typing.Union[Literal[3], Literal[4]]]]  # Error

crates/ruff/src/checkers/ast/mod.rs

@@ -2189,6 +2189,22 @@ where
                     }
                 }
 
+                // Ex) Union[...]
+                if self.enabled(Rule::UnnecessaryLiteralUnion) {
+                    let mut check = true;
+
+                    // Avoid duplicate checks if the parent is an `Union[...]`
+                    if let Some(Expr::Subscript(ast::ExprSubscript { value, .. })) =
+                        self.semantic.expr_grandparent()
+                    {
+                        check = !self.semantic.match_typing_expr(value, "Union");
+                    }
+
+                    if check {
+                        flake8_pyi::rules::unnecessary_literal_union(self, expr);
+                    }
+                }
+
                 if self.semantic.match_typing_expr(value, "Literal") {
                     self.semantic.flags |= SemanticModelFlags::LITERAL;
                 }
@@ -3136,18 +3152,24 @@ where
                 if self.is_stub {
                     if self.enabled(Rule::DuplicateUnionMember)
                         && self.semantic.in_type_definition()
-                        && self.semantic.expr_parent().map_or(true, |parent| {
-                            !matches!(
-                                parent,
-                                Expr::BinOp(ast::ExprBinOp {
-                                    op: Operator::BitOr,
-                                    ..
-                                })
-                            )
-                        })
+                        // Avoid duplicate checks if the parent is an `|`
+                        && !matches!(
+                            self.semantic.expr_parent(),
+                            Some(Expr::BinOp(ast::ExprBinOp { op: Operator::BitOr, ..}))
+                        )
                     {
                         flake8_pyi::rules::duplicate_union_member(self, expr);
                     }
+
+                    if self.enabled(Rule::UnnecessaryLiteralUnion)
+                        // Avoid duplicate checks if the parent is an `|`
+                        && !matches!(
+                            self.semantic.expr_parent(),
+                            Some(Expr::BinOp(ast::ExprBinOp { op: Operator::BitOr, ..}))
+                        )
+                    {
+                        flake8_pyi::rules::unnecessary_literal_union(self, expr);
+                    }
                 }
             }
             Expr::UnaryOp(ast::ExprUnaryOp {

crates/ruff/src/codes.rs

@@ -630,6 +630,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
         (Flake8Pyi, "024") => (RuleGroup::Unspecified, rules::flake8_pyi::rules::CollectionsNamedTuple),
         (Flake8Pyi, "025") => (RuleGroup::Unspecified, rules::flake8_pyi::rules::UnaliasedCollectionsAbcSetImport),
         (Flake8Pyi, "029") => (RuleGroup::Unspecified, rules::flake8_pyi::rules::StrOrReprDefinedInStub),
+        (Flake8Pyi, "030") => (RuleGroup::Unspecified, rules::flake8_pyi::rules::UnnecessaryLiteralUnion),
         (Flake8Pyi, "032") => (RuleGroup::Unspecified, rules::flake8_pyi::rules::AnyEqNeAnnotation),
         (Flake8Pyi, "033") => (RuleGroup::Unspecified, rules::flake8_pyi::rules::TypeCommentInStub),
         (Flake8Pyi, "034") => (RuleGroup::Unspecified, rules::flake8_pyi::rules::NonSelfReturnType),

crates/ruff/src/rules/flake8_pyi/mod.rs

@@ -52,6 +52,8 @@ mod tests {
     #[test_case(Rule::UnassignedSpecialVariableInStub, Path::new("PYI035.pyi"))]
     #[test_case(Rule::StrOrReprDefinedInStub, Path::new("PYI029.py"))]
     #[test_case(Rule::StrOrReprDefinedInStub, Path::new("PYI029.pyi"))]
+    #[test_case(Rule::UnnecessaryLiteralUnion, Path::new("PYI030.py"))]
+    #[test_case(Rule::UnnecessaryLiteralUnion, Path::new("PYI030.pyi"))]
     #[test_case(Rule::StubBodyMultipleStatements, Path::new("PYI048.py"))]
     #[test_case(Rule::StubBodyMultipleStatements, Path::new("PYI048.pyi"))]
     #[test_case(Rule::TSuffixedTypeAlias, Path::new("PYI043.py"))]

crates/ruff/src/rules/flake8_pyi/rules/mod.rs

@@ -22,6 +22,7 @@ pub(crate) use stub_body_multiple_statements::*;
 pub(crate) use type_alias_naming::*;
 pub(crate) use type_comment_in_stub::*;
 pub(crate) use unaliased_collections_abc_set_import::*;
+pub(crate) use unnecessary_literal_union::*;
 pub(crate) use unrecognized_platform::*;
 pub(crate) use unrecognized_version_info::*;
 
@@ -49,5 +50,6 @@ mod stub_body_multiple_statements;
 mod type_alias_naming;
 mod type_comment_in_stub;
 mod unaliased_collections_abc_set_import;
+mod unnecessary_literal_union;
 mod unrecognized_platform;
 mod unrecognized_version_info;

crates/ruff/src/rules/flake8_pyi/rules/unnecessary_literal_union.rs

@@ -0,0 +1,120 @@
+use ruff_python_semantic::SemanticModel;
+use rustpython_parser::ast::{self, Expr, Operator, Ranged};
+
+use ruff_diagnostics::{Diagnostic, Violation};
+use ruff_macros::{derive_message_formats, violation};
+use smallvec::SmallVec;
+
+use crate::checkers::ast::Checker;
+
+/// ## What it does
+/// Checks for the presence of multiple literal types in a union.
+///
+/// ## Why is this bad?
+/// Literal types accept multiple arguments and it is clearer to specify them
+/// as a single literal.
+///
+/// ## Example
+/// ```python
+/// from typing import Literal
+///
+/// field: Literal[1] | Literal[2]
+/// ```
+///
+/// Use instead:
+/// ```python
+/// from typing import Literal
+///
+/// field: Literal[1, 2]
+/// ```
+#[violation]
+pub struct UnnecessaryLiteralUnion {
+    members: Vec<String>,
+}
+
+impl Violation for UnnecessaryLiteralUnion {
+    #[derive_message_formats]
+    fn message(&self) -> String {
+        format!(
+            "Multiple literal members in a union. Use a single literal, e.g. `Literal[{}]`",
+            self.members.join(", ")
+        )
+    }
+}
+
+/// PYI030
+pub(crate) fn unnecessary_literal_union<'a>(checker: &mut Checker, expr: &'a Expr) {
+    let mut literal_exprs = SmallVec::<[&Box<Expr>; 1]>::new();
+
+    // Adds a member to `literal_exprs` if it is a `Literal` annotation
+    let mut collect_literal_expr = |expr: &'a Expr| {
+        if let Expr::Subscript(ast::ExprSubscript { value, slice, .. }) = expr {
+            if checker.semantic().match_typing_expr(value, "Literal") {
+                literal_exprs.push(slice);
+            }
+        }
+    };
+
+    // Traverse the union, collect all literal members
+    traverse_union(&mut collect_literal_expr, expr, checker.semantic());
+
+    // Raise a violation if more than one
+    if literal_exprs.len() > 1 {
+        let diagnostic = Diagnostic::new(
+            UnnecessaryLiteralUnion {
+                members: literal_exprs
+                    .into_iter()
+                    .map(|literal_expr| checker.locator.slice(literal_expr.range()).to_string())
+                    .collect(),
+            },
+            expr.range(),
+        );
+
+        checker.diagnostics.push(diagnostic);
+    }
+}
+
+/// Traverse a "union" type annotation, calling `func` on each expression in the union.
+fn traverse_union<'a, F>(func: &mut F, expr: &'a Expr, semantic: &SemanticModel)
+where
+    F: FnMut(&'a Expr),
+{
+    // Ex) x | y
+    if let Expr::BinOp(ast::ExprBinOp {
+        op: Operator::BitOr,
+        left,
+        right,
+        range: _,
+    }) = expr
+    {
+        // The union data structure usually looks like this:
+        //  a | b | c -> (a | b) | c
+        //
+        // However, parenthesized expressions can coerce it into any structure:
+        //  a | (b | c)
+        //
+        // So we have to traverse both branches in order (left, then right), to report members
+        // in the order they appear in the source code.
+
+        // Traverse the left then right arms
+        traverse_union(func, left, semantic);
+        traverse_union(func, right, semantic);
+        return;
+    }
+
+    // Ex) Union[x, y]
+    if let Expr::Subscript(ast::ExprSubscript { value, slice, .. }) = expr {
+        if semantic.match_typing_expr(value, "Union") {
+            if let Expr::Tuple(ast::ExprTuple { elts, .. }) = slice.as_ref() {
+                // Traverse each element of the tuple within the union recursively to handle cases
+                // such as `Union[..., Union[...]]
+                elts.iter()
+                    .for_each(|elt| traverse_union(func, elt, semantic));
+                return;
+            }
+        }
+    }
+
+    // Otherwise, call the function on expression
+    func(expr);
+}

crates/ruff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI030_PYI030.py.snap

@@ -0,0 +1,4 @@
+---
+source: crates/ruff/src/rules/flake8_pyi/mod.rs
+---
+