[flake8-pyi, ruff] Fix traversal of nested literals and unions (PYI016, PYI051, PYI055, PYI062, RUF041) (#14641)

crates/ruff_python_semantic/src/analyze/typing.rs

@@ -373,7 +373,7 @@ where
     ) where
         F: FnMut(&'a Expr, &'a Expr),
     {
-        // Ex) x | y
+        // Ex) `x | y`
         if let Expr::BinOp(ast::ExprBinOp {
             op: Operator::BitOr,
             left,
@@ -396,17 +396,21 @@ where
             return;
         }
 
-        // Ex) Union[x, y]
+        // Ex) `Union[x, y]`
         if let Expr::Subscript(ast::ExprSubscript { value, slice, .. }) = expr {
             if semantic.match_typing_expr(value, "Union") {
                 if let Expr::Tuple(tuple) = &**slice {
                     // Traverse each element of the tuple within the union recursively to handle cases
-                    // such as `Union[..., Union[...]]
+                    // such as `Union[..., Union[...]]`
                     tuple
                         .iter()
                         .for_each(|elem| inner(func, semantic, elem, Some(expr)));
                     return;
                 }
+
+                // Ex) `Union[Union[a, b]]` and `Union[a | b | c]`
+                inner(func, semantic, slice, Some(expr));
+                return;
             }
         }
 
@@ -435,18 +439,19 @@ where
     ) where
         F: FnMut(&'a Expr, &'a Expr),
     {
-        // Ex) Literal[x, y]
+        // Ex) `Literal[x, y]`
         if let Expr::Subscript(ast::ExprSubscript { value, slice, .. }) = expr {
             if semantic.match_typing_expr(value, "Literal") {
                 match &**slice {
                     Expr::Tuple(tuple) => {
                         // Traverse each element of the tuple within the literal recursively to handle cases
-                        // such as `Literal[..., Literal[...]]
+                        // such as `Literal[..., Literal[...]]`
                         for element in tuple {
                             inner(func, semantic, element, Some(expr));
                         }
                     }
                     other => {
+                        // Ex) `Literal[Literal[...]]`
                         inner(func, semantic, other, Some(expr));
                     }
                 }

crates/ruff_python_semantic/src/model.rs

@@ -5,7 +5,7 @@ use rustc_hash::FxHashMap;
 
 use ruff_python_ast::helpers::from_relative_import;
 use ruff_python_ast::name::{QualifiedName, UnqualifiedName};
-use ruff_python_ast::{self as ast, Expr, ExprContext, Operator, PySourceType, Stmt};
+use ruff_python_ast::{self as ast, Expr, ExprContext, PySourceType, Stmt};
 use ruff_text_size::{Ranged, TextRange, TextSize};
 
 use crate::binding::{
@@ -1506,38 +1506,48 @@ impl<'a> SemanticModel<'a> {
     /// Return `true` if the model is in a nested union expression (e.g., the inner `Union` in
     /// `Union[Union[int, str], float]`).
     pub fn in_nested_union(&self) -> bool {
-        // Ex) `Union[Union[int, str], float]`
-        if self
-            .current_expression_grandparent()
-            .and_then(Expr::as_subscript_expr)
-            .is_some_and(|parent| self.match_typing_expr(&parent.value, "Union"))
-        {
-            return true;
-        }
+        let mut parent_expressions = self.current_expressions().skip(1);
 
-        // Ex) `int | Union[str, float]`
-        if self.current_expression_parent().is_some_and(|parent| {
-            matches!(
-                parent,
-                Expr::BinOp(ast::ExprBinOp {
-                    op: Operator::BitOr,
-                    ..
-                })
-            )
-        }) {
-            return true;
+        match parent_expressions.next() {
+            // The parent expression is of the inner union is a single `typing.Union`.
+            // Ex) `Union[Union[a, b]]`
+            Some(Expr::Subscript(parent)) => self.match_typing_expr(&parent.value, "Union"),
+            // The parent expression is of the inner union is a tuple with two or more
+            // comma-separated elements and the parent of that tuple is a `typing.Union`.
+            // Ex) `Union[Union[a, b], Union[c, d]]`
+            Some(Expr::Tuple(_)) => parent_expressions
+                .next()
+                .and_then(Expr::as_subscript_expr)
+                .is_some_and(|grandparent| self.match_typing_expr(&grandparent.value, "Union")),
+            // The parent expression of the inner union is a PEP604-style union.
+            // Ex) `a | b | c` or `Union[a, b] | c`
+            // In contrast to `typing.Union`, PEP604-style unions are always binary operations, e.g.
+            // the expression `a | b | c` is represented by two binary unions: `(a | b) | c`.
+            Some(Expr::BinOp(bin_op)) => bin_op.op.is_bit_or(),
+            // Not a nested union otherwise.
+            _ => false,
         }
-
-        false
     }
 
     /// Return `true` if the model is in a nested literal expression (e.g., the inner `Literal` in
     /// `Literal[Literal[int, str], float]`).
     pub fn in_nested_literal(&self) -> bool {
-        // Ex) `Literal[Literal[int, str], float]`
-        self.current_expression_grandparent()
-            .and_then(Expr::as_subscript_expr)
-            .is_some_and(|parent| self.match_typing_expr(&parent.value, "Literal"))
+        let mut parent_expressions = self.current_expressions().skip(1);
+
+        match parent_expressions.next() {
+            // The parent expression of the current `Literal` is a tuple, and the
+            // grandparent is a `Literal`.
+            // Ex) `Literal[Literal[str], Literal[int]]`
+            Some(Expr::Tuple(_)) => parent_expressions
+                .next()
+                .and_then(Expr::as_subscript_expr)
+                .is_some_and(|grandparent| self.match_typing_expr(&grandparent.value, "Literal")),
+            // The parent expression of the current `Literal` is also a `Literal`.
+            // Ex) `Literal[Literal[str]]`
+            Some(Expr::Subscript(parent)) => self.match_typing_expr(&parent.value, "Literal"),
+            // Not a nested literal otherwise
+            _ => false,
+        }
     }
 
     /// Returns `true` if `left` and `right` are in the same branches of an `if`, `match`, or