Add PYI011 auto-fix

astral-sh · Mar 13, 2023 · f08fa77 · f08fa77
1 parent 375039e
commit f08fa77
Show file tree

Hide file tree

Showing 3 changed files with 205 additions and 56 deletions.
diff --git a/crates/ruff/resources/test/fixtures/flake8_pyi/PYI011.pyi b/crates/ruff/resources/test/fixtures/flake8_pyi/PYI011.pyi
@@ -109,5 +109,8 @@ def f35(
     + 1j,
 ) -> None: ...
 def f36(
-    x: str = sys.version,  # OK
+    *, x: str = sys.version,  # OK
+) -> None: ...
+def f37(
+    *, x: str = "" + "",  # Error PYI011 Only simple default values allowed for typed arguments
 ) -> None: ...
diff --git a/crates/ruff/src/rules/flake8_pyi/rules/simple_defaults.rs b/crates/ruff/src/rules/flake8_pyi/rules/simple_defaults.rs
@@ -1,20 +1,25 @@
 use rustpython_parser::ast::{Arguments, Constant, Expr, ExprKind, Operator, Unaryop};
 
-use ruff_diagnostics::{Diagnostic, Violation};
+use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Fix, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::types::Range;
 
 use crate::checkers::ast::Checker;
+use crate::registry::AsRule;
 
 #[violation]
 pub struct TypedArgumentSimpleDefaults;
 
 /// PYI011
-impl Violation for TypedArgumentSimpleDefaults {
+impl AlwaysAutofixableViolation for TypedArgumentSimpleDefaults {
     #[derive_message_formats]
     fn message(&self) -> String {
         format!("Only simple default values allowed for typed arguments")
     }
+
+    fn autofix_title(&self) -> String {
+        "Replace default value by `...`".to_string()
+    }
 }
 
 #[violation]
@@ -188,10 +193,18 @@ pub fn typed_argument_simple_defaults(checker: &mut Checker, args: &Arguments) {
             {
                 if arg.node.annotation.is_some() {
                     if !is_valid_default_value_with_annotation(default, checker) {
-                        checker.diagnostics.push(Diagnostic::new(
-                            TypedArgumentSimpleDefaults,
-                            Range::from(default),
-                        ));
+                        let mut diagnostic =
+                            Diagnostic::new(TypedArgumentSimpleDefaults, Range::from(default));
+
+                        if checker.patch(diagnostic.kind.rule()) {
+                            diagnostic.amend(Fix::replacement(
+                                "...".to_string(),
+                                default.location,
+                                default.end_location.unwrap(),
+                            ));
+                        }
+
+                        checker.diagnostics.push(diagnostic);
                     }
                 }
             }
@@ -207,10 +220,18 @@ pub fn typed_argument_simple_defaults(checker: &mut Checker, args: &Arguments) {
             {
                 if kwarg.node.annotation.is_some() {
                     if !is_valid_default_value_with_annotation(default, checker) {
-                        checker.diagnostics.push(Diagnostic::new(
-                            TypedArgumentSimpleDefaults,
-                            Range::from(default),
-                        ));
+                        let mut diagnostic =
+                            Diagnostic::new(TypedArgumentSimpleDefaults, Range::from(default));
+
+                        if checker.patch(diagnostic.kind.rule()) {
+                            diagnostic.amend(Fix::replacement(
+                                "...".to_string(),
+                                default.location,
+                                default.end_location.unwrap(),
+                            ));
+                        }
+
+                        checker.diagnostics.push(diagnostic);
                     }
                 }
             }