[flake8-pyi] Implement autofix and handle nested unions with single element (PYI041 , PYI055)

crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI016.py

@@ -84,3 +84,27 @@ def func2() -> str | str:  # PYI016: Duplicate union member `str`
 # duplicates of the outer `int`), but not three times (which would indicate that
 # we incorrectly re-checked the nested union).
 field25: typing.Union[int, int | int]  # PYI016: Duplicate union member `int`
+
+# Should emit in cases with nested `typing.Union`
+field26: typing.Union[typing.Union[int, int]]  # PYI016: Duplicate union member `int`
+
+# Should emit in cases with nested `typing.Union`
+field27: typing.Union[typing.Union[typing.Union[int, int]]]  # PYI016: Duplicate union member `int`
+
+# Should emit in cases with mixed `typing.Union` and `|`
+field28: typing.Union[int | int]  # Error
+
+# Should emit twice in cases with multiple nested `typing.Union`
+field29: typing.Union[int, typing.Union[typing.Union[int, int]]]  # Error
+
+# Should emit once in cases with multiple nested `typing.Union`
+field30: typing.Union[int, typing.Union[typing.Union[int, str]]]  # Error
+
+# Should emit once, and fix to `typing.Union[float, int]`
+field31: typing.Union[float, typing.Union[int | int]]  # Error
+
+# Should emit once, and fix to `typing.Union[float, int]`
+field32: typing.Union[float, typing.Union[int | int | int]]  # Error
+
+# Test case for mixed union type fix
+field33: typing.Union[typing.Union[int | int] | typing.Union[int | int]] # Error

crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI016.pyi

@@ -84,3 +84,27 @@ field24: typing.Union[int, typing.Union[int, int]]  # PYI016: Duplicate union me
 # duplicates of the outer `int`), but not three times (which would indicate that
 # we incorrectly re-checked the nested union).
 field25: typing.Union[int, int | int]  # PYI016: Duplicate union member `int`
+
+# Should emit in cases with nested `typing.Union`
+field26: typing.Union[typing.Union[int, int]]  # PYI016: Duplicate union member `int`
+
+# Should emit in cases with nested `typing.Union`
+field27: typing.Union[typing.Union[typing.Union[int, int]]]  # PYI016: Duplicate union member `int`
+
+# Should emit in cases with mixed `typing.Union` and `|`
+field28: typing.Union[int | int]  # Error
+
+# Should emit twice in cases with multiple nested `typing.Union`
+field29: typing.Union[int, typing.Union[typing.Union[int, int]]]  # Error
+
+# Should emit once in cases with multiple nested `typing.Union`
+field30: typing.Union[int, typing.Union[typing.Union[int, str]]]  # Error
+
+# Should emit once, and fix to `typing.Union[float, int]`
+field31: typing.Union[float, typing.Union[int | int]]  # Error
+
+# Should emit once, and fix to `typing.Union[float, int]`
+field32: typing.Union[float, typing.Union[int | int | int]]  # Error
+
+# Test case for mixed union type fix
+field33: typing.Union[typing.Union[int | int] | typing.Union[int | int]] # Error

crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI041.py

@@ -39,9 +39,55 @@ async def f4(**kwargs: int | int | float) -> None:
     ...
 
 
+
+def f5(arg1: int, *args: Union[int, int, float]) -> None: 
+    ...
+
+
+def f6(arg1: int, *args: Union[Union[int, int, float]]) -> None: 
+    ...
+
+
+def f7(arg1: int, *args: Union[Union[Union[int, int, float]]]) -> None: 
+    ...
+
+
+def f8(arg1: int, *args: Union[Union[Union[int | int | float]]]) -> None: 
+    ...
+
+
+def f9(
+    arg: Union[  # comment 
+        float, # another
+        complex, int]
+    ) -> None: 
+    ...
+
+def f10(
+    arg: (
+        int | # comment
+        float |  # another
+        complex
+    )    
+    ) -> None: 
+    ...
+
+
 class Foo:
     def good(self, arg: int) -> None:
         ...
 
     def bad(self, arg: int | float | complex) -> None:
         ...
+
+    def bad2(self, arg: int | Union[float, complex]) -> None: 
+        ...
+
+    def bad3(self, arg: Union[Union[float, complex], int]) -> None: 
+        ...
+
+    def bad4(self, arg: Union[float | complex, int]) -> None: 
+        ...
+
+    def bad5(self, arg: int | (float | complex)) -> None: 
+        ...

crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI041.pyi

@@ -33,7 +33,41 @@ def f3(arg1: int, *args: Union[int | int | float]) -> None: ...  # PYI041
 async def f4(**kwargs: int | int | float) -> None: ...  # PYI041
 
 
+def f5(arg1: int, *args: Union[int, int, float]) -> None: ...  # PYI041
+
+
+def f6(arg1: int, *args: Union[Union[int, int, float]]) -> None: ...  # PYI041
+
+
+def f7(arg1: int, *args: Union[Union[Union[int, int, float]]]) -> None: ...  # PYI041
+
+
+def f8(arg1: int, *args: Union[Union[Union[int | int | float]]]) -> None: ...  # PYI041
+
+
+def f9(
+    arg: Union[  # comment 
+        float, # another
+        complex, int]
+    ) -> None: ...  # PYI041
+
+def f10(
+    arg: (
+        int | # comment
+        float |  # another
+        complex
+    )    
+    ) -> None: ... # PYI041
+
 class Foo:
     def good(self, arg: int) -> None: ...
 
     def bad(self, arg: int | float | complex) -> None: ...  # PYI041
+
+    def bad2(self, arg: int | Union[float, complex]) -> None: ...  # PYI041
+
+    def bad3(self, arg: Union[Union[float, complex], int]) -> None: ...  # PYI041
+
+    def bad4(self, arg: Union[float | complex, int]) -> None: ...  # PYI041
+
+    def bad5(self, arg: int | (float | complex)) -> None: ...  # PYI041

crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI051.py

@@ -5,6 +5,9 @@
 B: TypeAlias = typing.Union[Literal[b"bar", b"foo"], bytes, str]
 C: TypeAlias = typing.Union[Literal[5], int, typing.Union[Literal["foo"], str]]
 D: TypeAlias = typing.Union[Literal[b"str_bytes", 42], bytes, int]
+E: TypeAlias = typing.Union[typing.Union[typing.Union[typing.Union[Literal["foo"], str]]]]
+F: TypeAlias = typing.Union[str, typing.Union[typing.Union[typing.Union[Literal["foo"], int]]]]
+G: typing.Union[str, typing.Union[typing.Union[typing.Union[Literal["foo"], int]]]]
 
 def func(x: complex | Literal[1J], y: Union[Literal[3.14], float]): ...
 

crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI051.pyi

@@ -5,6 +5,9 @@ A: str | Literal["foo"]
 B: TypeAlias = typing.Union[Literal[b"bar", b"foo"], bytes, str]
 C: TypeAlias = typing.Union[Literal[5], int, typing.Union[Literal["foo"], str]]
 D: TypeAlias = typing.Union[Literal[b"str_bytes", 42], bytes, int]
+E: TypeAlias = typing.Union[typing.Union[typing.Union[typing.Union[Literal["foo"], str]]]]
+F: TypeAlias = typing.Union[str, typing.Union[typing.Union[typing.Union[Literal["foo"], int]]]]
+G: typing.Union[str, typing.Union[typing.Union[typing.Union[Literal["foo"], int]]]]
 
 def func(x: complex | Literal[1J], y: Union[Literal[3.14], float]): ...
 

crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI055.py

@@ -1,11 +1,14 @@
 import builtins
 from typing import Union
 
-w: builtins.type[int] | builtins.type[str] | builtins.type[complex]
-x: type[int] | type[str] | type[float]
-y: builtins.type[int] | type[str] | builtins.type[complex]
-z: Union[type[float], type[complex]]
-z: Union[type[float, int], type[complex]]
+s: builtins.type[int] | builtins.type[str] | builtins.type[complex]
+t: type[int] | type[str] | type[float]
+u: builtins.type[int] | type[str] | builtins.type[complex]
+v: Union[type[float], type[complex]]
+w: Union[type[float, int], type[complex]]
+x: Union[Union[type[float, int], type[complex]]]
+y: Union[Union[Union[type[float, int], type[complex]]]]
+z: Union[type[complex], Union[Union[type[float, int]]]]
 
 
 def func(arg: type[int] | str | type[float]) -> None:
@@ -30,6 +33,9 @@ def func():
     # PYI055
     x: type[requests_mock.Mocker] | type[httpretty] | type[str] = requests_mock.Mocker
     y: Union[type[requests_mock.Mocker], type[httpretty], type[str]] = requests_mock.Mocker
+    z: Union[  # comment
+        type[requests_mock.Mocker], # another comment
+        type[httpretty], type[str]] = requests_mock.Mocker
 
 
 def func():

crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI055.pyi

@@ -1,11 +1,14 @@
 import builtins
 from typing import Union
 
-w: builtins.type[int] | builtins.type[str] | builtins.type[complex]
-x: type[int] | type[str] | type[float]
-y: builtins.type[int] | type[str] | builtins.type[complex]
-z: Union[type[float], type[complex]]
-z: Union[type[float, int], type[complex]]
+s: builtins.type[int] | builtins.type[str] | builtins.type[complex]
+t: type[int] | type[str] | type[float]
+u: builtins.type[int] | type[str] | builtins.type[complex]
+v: Union[type[float], type[complex]]
+w: Union[type[float, int], type[complex]]
+x: Union[Union[type[float, int], type[complex]]]
+y: Union[Union[Union[type[float, int], type[complex]]]]
+z: Union[type[complex], Union[Union[type[float, int]]]]
 
 def func(arg: type[int] | str | type[float]) -> None: ...
 
@@ -16,10 +19,13 @@ z: Union[float, complex]
 
 def func(arg: type[int, float] | str) -> None: ...
 
-# OK
+# PYI055
 item: type[requests_mock.Mocker] | type[httpretty] = requests_mock.Mocker
 
 def func():
     # PYI055
     item: type[requests_mock.Mocker] | type[httpretty] | type[str] = requests_mock.Mocker
     item2: Union[type[requests_mock.Mocker], type[httpretty], type[str]] = requests_mock.Mocker
+    item3: Union[  # comment
+        type[requests_mock.Mocker], # another comment
+        type[httpretty], type[str]] = requests_mock.Mocker

crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI062.py

@@ -14,8 +14,20 @@
 Literal[1, Literal[2], Literal[2]]  # once
 t.Literal[1, t.Literal[2, t.Literal[1]]]  # once
 typing_extensions.Literal[1, 1, 1]  # twice
+Literal[
+    1, # comment
+    Literal[ # another comment
+        1
+    ]
+]  # once
 
 # Ensure issue is only raised once, even on nested literals
 MyType = Literal["foo", Literal[True, False, True], "bar"]  # PYI062
 
 n: Literal["No", "duplicates", "here", 1, "1"]
+
+
+# nested literals, all equivalent to `Literal[1]`
+Literal[Literal[1]]  # no duplicate
+Literal[Literal[Literal[1], Literal[1]]]  # once
+Literal[Literal[1], Literal[Literal[Literal[1]]]]  # once

crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI062.pyi

@@ -2,20 +2,32 @@ from typing import Literal
 import typing as t
 import typing_extensions
 
-x: Literal[True, False, True, False]  # PY062 twice here
+x: Literal[True, False, True, False]  # PYI062 twice here
 
-y: Literal[1, print("hello"), 3, Literal[4, 1]]  # PY062 on the last 1
+y: Literal[1, print("hello"), 3, Literal[4, 1]]  # PYI062 on the last 1
 
-z: Literal[{1, 3, 5}, "foobar", {1,3,5}]  # PY062 on the set literal
+z: Literal[{1, 3, 5}, "foobar", {1,3,5}]  # PYI062 on the set literal
 
 Literal[1, Literal[1]]  # once
 Literal[1, 2, Literal[1, 2]]  # twice
 Literal[1, Literal[1], Literal[1]]  # twice
 Literal[1, Literal[2], Literal[2]]  # once
 t.Literal[1, t.Literal[2, t.Literal[1]]]  # once
 typing_extensions.Literal[1, 1, 1]  # twice
+Literal[
+    1, # comment
+    Literal[ # another comment
+        1
+    ]
+]  # once
 
 # Ensure issue is only raised once, even on nested literals
 MyType = Literal["foo", Literal[True, False, True], "bar"]  # PYI062
 
 n: Literal["No", "duplicates", "here", 1, "1"]
+
+
+# nested literals, all equivalent to `Literal[1]`
+Literal[Literal[1]]  # no duplicate
+Literal[Literal[Literal[1], Literal[1]]]  # once
+Literal[Literal[1], Literal[Literal[Literal[1]]]]  # once

crates/ruff_linter/src/rules/flake8_pyi/rules/duplicate_literal_member.rs

@@ -2,7 +2,7 @@ use std::collections::HashSet;
 
 use rustc_hash::FxHashSet;
 
-use ruff_diagnostics::{AlwaysFixableViolation, Diagnostic, Edit, Fix};
+use ruff_diagnostics::{AlwaysFixableViolation, Applicability, Diagnostic, Edit, Fix};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::comparable::ComparableExpr;
 use ruff_python_ast::{self as ast, Expr, ExprContext};
@@ -73,33 +73,39 @@ pub(crate) fn duplicate_literal_member<'a>(checker: &mut Checker, expr: &'a Expr
     // Traverse the literal, collect all diagnostic members.
     traverse_literal(&mut check_for_duplicate_members, checker.semantic(), expr);
 
+    if diagnostics.is_empty() {
+        return;
+    }
+
     // If there's at least one diagnostic, create a fix to remove the duplicate members.
-    if !diagnostics.is_empty() {
-        if let Expr::Subscript(subscript) = expr {
-            let subscript = Expr::Subscript(ast::ExprSubscript {
-                slice: Box::new(if let [elt] = unique_nodes.as_slice() {
-                    (*elt).clone()
-                } else {
-                    Expr::Tuple(ast::ExprTuple {
-                        elts: unique_nodes.into_iter().cloned().collect(),
-                        range: TextRange::default(),
-                        ctx: ExprContext::Load,
-                        parenthesized: false,
-                    })
-                }),
-                value: subscript.value.clone(),
-                range: TextRange::default(),
-                ctx: ExprContext::Load,
-            });
-            let fix = Fix::safe_edit(Edit::range_replacement(
-                checker.generator().expr(&subscript),
-                expr.range(),
-            ));
-            for diagnostic in &mut diagnostics {
-                diagnostic.set_fix(fix.clone());
-            }
+    if let Expr::Subscript(subscript) = expr {
+        let subscript = Expr::Subscript(ast::ExprSubscript {
+            slice: Box::new(if let [elt] = unique_nodes.as_slice() {
+                (*elt).clone()
+            } else {
+                Expr::Tuple(ast::ExprTuple {
+                    elts: unique_nodes.into_iter().cloned().collect(),
+                    range: TextRange::default(),
+                    ctx: ExprContext::Load,
+                    parenthesized: false,
+                })
+            }),
+            value: subscript.value.clone(),
+            range: TextRange::default(),
+            ctx: ExprContext::Load,
+        });
+        let fix = Fix::applicable_edit(
+            Edit::range_replacement(checker.generator().expr(&subscript), expr.range()),
+            if checker.comment_ranges().intersects(expr.range()) {
+                Applicability::Unsafe
+            } else {
+                Applicability::Safe
+            },
+        );
+        for diagnostic in &mut diagnostics {
+            diagnostic.set_fix(fix.clone());
         }
-    }
+    };
 
     checker.diagnostics.append(&mut diagnostics);
 }