Fix PYI011 and add auto-fix (#3492)

astral-sh · Mar 14, 2023 · 7e90411 · 7e90411
1 parent 344daeb
commit 7e90411
Show file tree

Hide file tree

Showing 3 changed files with 242 additions and 135 deletions.
diff --git a/crates/ruff/resources/test/fixtures/flake8_pyi/PYI011.pyi b/crates/ruff/resources/test/fixtures/flake8_pyi/PYI011.pyi
@@ -1,3 +1,10 @@
+import math
+import os
+import sys
+from math import inf
+
+import numpy as np
+
 def f12(
     x,
     y: str = os.pathsep,  # Error PYI011 Only simple default values allowed for typed arguments
@@ -101,3 +108,9 @@ def f35(
     x: complex = math.inf  # Error PYI011 Only simple default values allowed for typed arguments
     + 1j,
 ) -> None: ...
+def f36(
+    *, 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]
@@ -111,7 +116,7 @@ fn is_valid_default_value_with_annotation(default: &Expr, checker: &Checker) ->
             if let ExprKind::Attribute { .. } = &operand.node {
                 if checker
                     .ctx
-                    .resolve_call_path(default)
+                    .resolve_call_path(operand)
                     .map_or(false, |call_path| {
                         ALLOWED_MATH_ATTRIBUTES_IN_DEFAULTS.iter().any(|target| {
                             // reject `-math.nan`
@@ -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);
                     }
                 }
             }