Avoid applying PEP 646 rewrites in invalid contexts

Author: charliermarsh

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

@@ -1,27 +1,58 @@
 from typing import Generic, TypeVarTuple, Unpack
 
-Shape = TypeVarTuple('Shape')
+Shape = TypeVarTuple("Shape")
+
 
 class C(Generic[Unpack[Shape]]):
     pass
 
-class D(Generic[Unpack  [Shape]]):
+
+class D(Generic[Unpack[Shape]]):
     pass
 
-def f(*args: Unpack[tuple[int, ...]]): pass
 
-def f(*args: Unpack[other.Type]): pass
+def f(*args: Unpack[tuple[int, ...]]):
+    pass
 
 
-# Not valid unpackings but they are valid syntax
-def foo(*args: Unpack[int | str]) -> None: pass
-def foo(*args: Unpack[int and str]) -> None: pass
-def foo(*args: Unpack[int > str]) -> None: pass
+def f(*args: Unpack[other.Type]):
+    pass
+
+
+def f(*args: Generic[int, Unpack[int]]):
+    pass
+
+
+# Valid syntax, but can't be unpacked.
+def f(*args: Unpack[int | str]) -> None:
+    pass
+
+
+def f(*args: Unpack[int and str]) -> None:
+    pass
+
+
+def f(*args: Unpack[int > str]) -> None:
+    pass
+
 
-# We do not use the shorthand unpacking syntax in the following cases
 from typing import TypedDict
+
+
 class KwargsDict(TypedDict):
     x: int
     y: int
 
-def foo(name: str, /, **kwargs: Unpack[KwargsDict]) -> None: pass
+
+# OK
+def f(name: str, /, **kwargs: Unpack[KwargsDict]) -> None:
+    pass
+
+
+# OK
+def f() -> object:
+    return Unpack[tuple[int, ...]]
+
+
+# OK
+def f(x: Unpack[int]) -> object: ...

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

@@ -145,11 +145,9 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
             if checker.enabled(Rule::FStringNumberFormat) {
                 refurb::rules::fstring_number_format(checker, subscript);
             }
-
             if checker.enabled(Rule::IncorrectlyParenthesizedTupleInSubscript) {
                 ruff::rules::subscript_with_parenthesized_tuple(checker, subscript);
             }
-
             if checker.enabled(Rule::NonPEP646Unpack) {
                 pyupgrade::rules::use_pep646_unpack(checker, subscript);
             }

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

@@ -1,6 +1,7 @@
 use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::ExprSubscript;
+use ruff_python_semantic::SemanticModel;
 
 use crate::{checkers::ast::Checker, settings::types::PythonVersion};
 
@@ -14,7 +15,6 @@ use crate::{checkers::ast::Checker, settings::types::PythonVersion};
 /// `Unpack[]` syntax.
 ///
 /// ## Example
-///
 /// ```python
 /// from typing import Unpack
 ///
@@ -24,12 +24,16 @@ use crate::{checkers::ast::Checker, settings::types::PythonVersion};
 /// ```
 ///
 /// Use instead:
-///
 /// ```python
 /// def foo(*args: *tuple[int, ...]) -> None:
 ///     pass
 /// ```
 ///
+/// ## Fix safety
+/// This rule's fix is marked as unsafe, as `Unpack[T]` and `*T` are considered
+/// different values when introspecting types at runtime. However, in most cases,
+/// the fix should be safe to apply.
+///
 /// [PEP 646]: https://peps.python.org/pep-0646/
 #[violation]
 pub struct NonPEP646Unpack;
@@ -57,47 +61,88 @@ pub(crate) fn use_pep646_unpack(checker: &mut Checker, expr: &ExprSubscript) {
         return;
     }
 
-    // Ignore `kwarg` unpacking, as in:
-    // ```python
-    // def f(**kwargs: Unpack[Array]):
-    //     ...
-    // ```
-    if checker
-        .semantic()
-        .current_statement()
-        .as_function_def_stmt()
-        .and_then(|stmt| stmt.parameters.kwarg.as_ref())
-        .and_then(|kwarg| kwarg.annotation.as_ref())
-        .and_then(|annotation| annotation.as_subscript_expr())
-        .is_some_and(|subscript| subscript == expr)
-    {
-        return;
-    }
-
     let ExprSubscript {
         range,
         value,
         slice,
         ..
     } = expr;
 
+    // Skip semantically invalid subscript calls (e.g. `Unpack[str | num]`).
+    if !(slice.is_name_expr() || slice.is_subscript_expr() || slice.is_attribute_expr()) {
+        return;
+    }
+
     if !checker.semantic().match_typing_expr(value, "Unpack") {
         return;
     }
 
-    // Skip semantically invalid subscript calls (e.g. `Unpack[str | num]`).
-    if !(slice.is_name_expr() || slice.is_subscript_expr() || slice.is_attribute_expr()) {
+    // Determine whether we're in a valid context for a star expression.
+    //
+    // Star expressions are only allowed in two places:
+    // - Subscript indexes (e.g., `Generic[DType, *Shape]`).
+    // - Variadic positional arguments (e.g., `def f(*args: *int)`).
+    //
+    // See: <https://peps.python.org/pep-0646/#grammar-changes>
+    if !in_subscript_index(expr, checker.semantic()) && !in_vararg(expr, checker.semantic()) {
         return;
     }
 
     let mut diagnostic = Diagnostic::new(NonPEP646Unpack, *range);
-
-    let inner = checker.locator().slice(slice.as_ref());
-
-    diagnostic.set_fix(Fix::safe_edit(Edit::range_replacement(
-        format!("*{inner}"),
+    diagnostic.set_fix(Fix::unsafe_edit(Edit::range_replacement(
+        format!("*{}", checker.locator().slice(slice.as_ref())),
         *range,
     )));
-
     checker.diagnostics.push(diagnostic);
 }
+
+/// Determine whether the [`ExprSubscript`] is in a subscript index (e.g., `Generic[Unpack[int]]`).
+fn in_subscript_index(expr: &ExprSubscript, semantic: &SemanticModel) -> bool {
+    let parent = semantic
+        .current_expressions()
+        .skip(1)
+        .find_map(|expr| expr.as_subscript_expr());
+
+    let Some(parent) = parent else {
+        return false;
+    };
+
+    // E.g., `Generic[Unpack[int]]`.
+    if parent
+        .slice
+        .as_subscript_expr()
+        .is_some_and(|slice| slice == expr)
+    {
+        return true;
+    }
+
+    // E.g., `Generic[DType, Unpack[int]]`.
+    if parent.slice.as_tuple_expr().map_or(false, |slice| {
+        slice
+            .elts
+            .iter()
+            .any(|elt| elt.as_subscript_expr().is_some_and(|elt| elt == expr))
+    }) {
+        return true;
+    }
+
+    false
+}
+
+/// Determine whether the [`ExprSubscript`] is attached to a variadic argument in a function
+/// definition (e.g., `def f(*args: Unpack[int])`).
+fn in_vararg(expr: &ExprSubscript, semantic: &SemanticModel) -> bool {
+    let parent = semantic.current_statement().as_function_def_stmt();
+
+    let Some(parent) = parent else {
+        return false;
+    };
+
+    parent
+        .parameters
+        .vararg
+        .as_ref()
+        .and_then(|vararg| vararg.annotation.as_ref())
+        .and_then(|annotation| annotation.as_subscript_expr())
+        .map_or(false, |annotation| annotation == expr)
+}

crates/ruff_linter/src/rules/pyupgrade/snapshots/ruff_linter__rules__pyupgrade__tests__unpack_pep_646_py311.snap

@@ -1,83 +1,92 @@
 ---
 source: crates/ruff_linter/src/rules/pyupgrade/mod.rs
-snapshot_kind: text
 ---
-UP044.py:5:17: UP044 [*] Use `*` for unpacking
+UP044.py:6:17: UP044 [*] Use `*` for unpacking
   |
-3 | Shape = TypeVarTuple('Shape')
-4 | 
-5 | class C(Generic[Unpack[Shape]]):
+6 | class C(Generic[Unpack[Shape]]):
   |                 ^^^^^^^^^^^^^ UP044
-6 |     pass
+7 |     pass
   |
   = help: Convert to `*` for unpacking
 
 ℹ Safe fix
-2 2 | 
-3 3 | Shape = TypeVarTuple('Shape')
+3 3 | Shape = TypeVarTuple("Shape")
 4 4 | 
-5   |-class C(Generic[Unpack[Shape]]):
-  5 |+class C(Generic[*Shape]):
-6 6 |     pass
-7 7 | 
-8 8 | class D(Generic[Unpack  [Shape]]):
+5 5 | 
+6   |-class C(Generic[Unpack[Shape]]):
+  6 |+class C(Generic[*Shape]):
+7 7 |     pass
+8 8 | 
+9 9 | 
 
-UP044.py:8:17: UP044 [*] Use `*` for unpacking
-  |
-6 |     pass
-7 | 
-8 | class D(Generic[Unpack  [Shape]]):
-  |                 ^^^^^^^^^^^^^^^ UP044
-9 |     pass
-  |
-  = help: Convert to `*` for unpacking
+UP044.py:10:17: UP044 [*] Use `*` for unpacking
+   |
+10 | class D(Generic[Unpack[Shape]]):
+   |                 ^^^^^^^^^^^^^ UP044
+11 |     pass
+   |
+   = help: Convert to `*` for unpacking
 
 ℹ Safe fix
-5 5 | class C(Generic[Unpack[Shape]]):
-6 6 |     pass
-7 7 | 
-8   |-class D(Generic[Unpack  [Shape]]):
-  8 |+class D(Generic[*Shape]):
-9 9 |     pass
-10 10 | 
-11 11 | def f(*args: Unpack[tuple[int, ...]]): pass
+7  7  |     pass
+8  8  | 
+9  9  | 
+10    |-class D(Generic[Unpack[Shape]]):
+   10 |+class D(Generic[*Shape]):
+11 11 |     pass
+12 12 | 
+13 13 | 
 
-UP044.py:11:14: UP044 [*] Use `*` for unpacking
+UP044.py:14:14: UP044 [*] Use `*` for unpacking
    |
- 9 |     pass
-10 | 
-11 | def f(*args: Unpack[tuple[int, ...]]): pass
+14 | def f(*args: Unpack[tuple[int, ...]]):
    |              ^^^^^^^^^^^^^^^^^^^^^^^ UP044
-12 | 
-13 | def f(*args: Unpack[other.Type]): pass
+15 |     pass
    |
    = help: Convert to `*` for unpacking
 
 ℹ Safe fix
-8  8  | class D(Generic[Unpack  [Shape]]):
-9  9  |     pass
-10 10 | 
-11    |-def f(*args: Unpack[tuple[int, ...]]): pass
-   11 |+def f(*args: *tuple[int, ...]): pass
+11 11 |     pass
 12 12 | 
-13 13 | def f(*args: Unpack[other.Type]): pass
-14 14 | 
+13 13 | 
+14    |-def f(*args: Unpack[tuple[int, ...]]):
+   14 |+def f(*args: *tuple[int, ...]):
+15 15 |     pass
+16 16 | 
+17 17 | 
 
-UP044.py:13:14: UP044 [*] Use `*` for unpacking
+UP044.py:18:14: UP044 [*] Use `*` for unpacking
    |
-11 | def f(*args: Unpack[tuple[int, ...]]): pass
-12 | 
-13 | def f(*args: Unpack[other.Type]): pass
+18 | def f(*args: Unpack[other.Type]):
    |              ^^^^^^^^^^^^^^^^^^ UP044
+19 |     pass
    |
    = help: Convert to `*` for unpacking
 
 ℹ Safe fix
-10 10 | 
-11 11 | def f(*args: Unpack[tuple[int, ...]]): pass
-12 12 | 
-13    |-def f(*args: Unpack[other.Type]): pass
-   13 |+def f(*args: *other.Type): pass
-14 14 | 
-15 15 | 
-16 16 | # Not valid unpackings but they are valid syntax
+15 15 |     pass
+16 16 | 
+17 17 | 
+18    |-def f(*args: Unpack[other.Type]):
+   18 |+def f(*args: *other.Type):
+19 19 |     pass
+20 20 | 
+21 21 | 
+
+UP044.py:22:27: UP044 [*] Use `*` for unpacking
+   |
+22 | def f(*args: Generic[int, Unpack[int]]):
+   |                           ^^^^^^^^^^^ UP044
+23 |     pass
+   |
+   = help: Convert to `*` for unpacking
+
+ℹ Safe fix
+19 19 |     pass
+20 20 | 
+21 21 | 
+22    |-def f(*args: Generic[int, Unpack[int]]):
+   22 |+def f(*args: Generic[int, *int]):
+23 23 |     pass
+24 24 | 
+25 25 |