Implement autofix for PYI001

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

@@ -1,11 +1,13 @@
+use ruff_text_size::TextSize;
 use std::fmt;
 
 use rustpython_parser::ast::{self, Expr, Ranged};
 
-use ruff_diagnostics::{Diagnostic, Violation};
+use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit, Fix};
 use ruff_macros::{derive_message_formats, violation};
 
 use crate::checkers::ast::Checker;
+use crate::registry::Rule;
 
 #[derive(Debug, PartialEq, Eq, Copy, Clone)]
 pub(crate) enum VarKind {
@@ -50,12 +52,17 @@ pub struct UnprefixedTypeParam {
     kind: VarKind,
 }
 
-impl Violation for UnprefixedTypeParam {
+impl AlwaysAutofixableViolation for UnprefixedTypeParam {
     #[derive_message_formats]
     fn message(&self) -> String {
         let UnprefixedTypeParam { kind } = self;
         format!("Name of private `{kind}` must start with `_`")
     }
+
+    fn autofix_title(&self) -> String {
+        let UnprefixedTypeParam { kind } = self;
+        format!("Prefix `{kind}` with `_`")
+    }
 }
 
 /// PYI001
@@ -69,7 +76,7 @@ pub(crate) fn prefix_type_params(checker: &mut Checker, value: &Expr, targets: &
         }
     };
 
-    if let Expr::Call(ast::ExprCall { func, .. }) = value {
+    if let Expr::Call(ast::ExprCall { func, args, .. }) = value {
         let Some(kind) = checker.semantic_model().resolve_call_path(func).and_then(|call_path| {
             if checker.semantic_model().match_typing_call_path(&call_path, "ParamSpec") {
                 Some(VarKind::ParamSpec)
@@ -83,8 +90,19 @@ pub(crate) fn prefix_type_params(checker: &mut Checker, value: &Expr, targets: &
         }) else {
             return;
         };
-        checker
-            .diagnostics
-            .push(Diagnostic::new(UnprefixedTypeParam { kind }, value.range()));
+        let mut diagnostic = Diagnostic::new(UnprefixedTypeParam { kind }, value.range());
+        if checker.patch(Rule::UnprefixedTypeParam) {
+            if let Expr::Name(ast::ExprName { range, .. }) = &targets[0] {
+                diagnostic.set_fix(Fix::automatic_edits(
+                    Edit::insertion(format!("_"), range.start()),
+                    [Edit::insertion(
+                        format!("_"),
+                        args[0].range().start() + TextSize::from(1),
+                    )],
+                ));
+            };
+        };
+
+        checker.diagnostics.push(diagnostic);
     }
 }

crates/ruff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI001_PYI001.pyi.snap

@@ -1,7 +1,7 @@
 ---
 source: crates/ruff/src/rules/flake8_pyi/mod.rs
 ---
-PYI001.pyi:3:5: PYI001 Name of private `TypeVar` must start with `_`
+PYI001.pyi:3:5: PYI001 [*] Name of private `TypeVar` must start with `_`
   |
 3 | from typing import ParamSpec, TypeVar, TypeVarTuple
 4 | 
@@ -10,8 +10,18 @@ PYI001.pyi:3:5: PYI001 Name of private `TypeVar` must start with `_`
 6 | 
 7 | TTuple = TypeVarTuple("TTuple")  # Error: TypeVarTuples must also start with _
   |
+  = help: Prefix `TypeVar` with `_`
 
-PYI001.pyi:5:10: PYI001 Name of private `TypeVarTuple` must start with `_`
+ℹ Fix
+1 1 | from typing import ParamSpec, TypeVar, TypeVarTuple
+2 2 | 
+3   |-T = TypeVar("T")  # Error: TypeVars in stubs must start with _
+  3 |+_T = TypeVar("_T")  # Error: TypeVars in stubs must start with _
+4 4 | 
+5 5 | TTuple = TypeVarTuple("TTuple")  # Error: TypeVarTuples must also start with _
+6 6 | 
+
+PYI001.pyi:5:10: PYI001 [*] Name of private `TypeVarTuple` must start with `_`
   |
 5 | T = TypeVar("T")  # Error: TypeVars in stubs must start with _
 6 | 
@@ -20,8 +30,19 @@ PYI001.pyi:5:10: PYI001 Name of private `TypeVarTuple` must start with `_`
 8 | 
 9 | P = ParamSpec("P")  # Error: ParamSpecs must start with _
   |
+  = help: Prefix `TypeVarTuple` with `_`
+
+ℹ Fix
+2 2 | 
+3 3 | T = TypeVar("T")  # Error: TypeVars in stubs must start with _
+4 4 | 
+5   |-TTuple = TypeVarTuple("TTuple")  # Error: TypeVarTuples must also start with _
+  5 |+_TTuple = TypeVarTuple("_TTuple")  # Error: TypeVarTuples must also start with _
+6 6 | 
+7 7 | P = ParamSpec("P")  # Error: ParamSpecs must start with _
+8 8 | 
 
-PYI001.pyi:7:5: PYI001 Name of private `ParamSpec` must start with `_`
+PYI001.pyi:7:5: PYI001 [*] Name of private `ParamSpec` must start with `_`
    |
  7 | TTuple = TypeVarTuple("TTuple")  # Error: TypeVarTuples must also start with _
  8 | 
@@ -30,5 +51,16 @@ PYI001.pyi:7:5: PYI001 Name of private `ParamSpec` must start with `_`
 10 | 
 11 | _T = TypeVar("_T")  # OK
    |
+   = help: Prefix `ParamSpec` with `_`
+
+ℹ Fix
+4 4 | 
+5 5 | TTuple = TypeVarTuple("TTuple")  # Error: TypeVarTuples must also start with _
+6 6 | 
+7   |-P = ParamSpec("P")  # Error: ParamSpecs must start with _
+  7 |+_P = ParamSpec("_P")  # Error: ParamSpecs must start with _
+8 8 | 
+9 9 | _T = TypeVar("_T")  # OK
+10 10 |