Implement autofix for TDO006

crates/ruff/src/checkers/tokens.rs

@@ -195,7 +195,7 @@ pub fn check_tokens(
     // T001, T002, T003, T004, T005, T006, T007
     if enforce_todos {
         diagnostics.extend(
-            flake8_todo::rules::check_rules(tokens)
+            flake8_todo::rules::check_todos(tokens, autofix, settings)
                 .into_iter()
                 .filter(|diagnostic| settings.rules.enabled(diagnostic.kind.rule())),
         );

crates/ruff/src/rules/flake8_todo/rules.rs

@@ -3,11 +3,16 @@ use std::iter::Peekable;
 use once_cell::sync::Lazy;
 
 use regex::{CaptureMatches, Regex, RegexSet};
-use ruff_diagnostics::{Diagnostic, Violation};
+use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit, Fix, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::types::Range;
-use rustpython_parser::lexer::LexResult;
 use rustpython_parser::Tok;
+use rustpython_parser::{ast::Location, lexer::LexResult};
+
+use crate::{
+    registry::Rule,
+    settings::{flags, Settings},
+};
 
 /// ## What it does
 /// Checks that a TODO comment is actually labelled with "TODO".
@@ -32,8 +37,6 @@ use rustpython_parser::Tok;
 pub struct InvalidTodoTag {
     pub tag: String,
 }
-
-// TODO - autofix this to just insert TODO instead of the tag?
 impl Violation for InvalidTodoTag {
     #[derive_message_formats]
     fn message(&self) -> String {
@@ -89,7 +92,7 @@ pub struct MissingLinkInTodo;
 impl Violation for MissingLinkInTodo {
     #[derive_message_formats]
     fn message(&self) -> String {
-        format!("Missing issue link in TODO")
+        format!("Missing issue link following TODO")
     }
 }
 
@@ -161,12 +164,16 @@ impl Violation for MissingTextInTodo {
 pub struct InvalidCapitalizationInTodo {
     pub tag: String,
 }
-impl Violation for InvalidCapitalizationInTodo {
+impl AlwaysAutofixableViolation for InvalidCapitalizationInTodo {
     #[derive_message_formats]
     fn message(&self) -> String {
         let InvalidCapitalizationInTodo { tag } = self;
         format!("Invalid TODO capitalization: `{tag}` should be `TODO`")
     }
+
+    fn autofix_title(&self) -> String {
+        "Fix capitalization in `TODO`".to_string()
+    }
 }
 
 /// ## What it does
@@ -232,7 +239,11 @@ static ISSUE_LINK_REGEX_SET: Lazy<RegexSet> = Lazy::new(|| {
 
 static NUM_CAPTURE_GROUPS: usize = 5usize;
 
-pub fn check_rules(tokens: &[LexResult]) -> Vec<Diagnostic> {
+pub fn check_todos(
+    tokens: &[LexResult],
+    autofix: flags::Autofix,
+    settings: &Settings,
+) -> Vec<Diagnostic> {
     let mut diagnostics: Vec<Diagnostic> = vec![];
     let mut prev_token_is_todo = false;
 
@@ -256,11 +267,10 @@ pub fn check_rules(tokens: &[LexResult]) -> Vec<Diagnostic> {
             diagnostics_ref.push(Diagnostic::new(MissingLinkInTodo, range));
         }
 
-        if let Some(captures_ref) = get_captured_matches(comment).peek() {
-            let captures = captures_ref.to_owned();
-
-            // captures.get(1) is the tag, which is required for the regex to match. The unwrap()
-            // call is therefore safe
+        // TODO: gotta be a way to remove this nest
+        if let Some(captures) = get_captured_matches(comment).peek() {
+            // Unwrap is safe because the "tag" capture group is required to enter the current
+            // block
             let tag = captures.name("tag").unwrap().as_str();
             if tag != "TODO" {
                 diagnostics_ref.push(Diagnostic::new(
@@ -271,12 +281,39 @@ pub fn check_rules(tokens: &[LexResult]) -> Vec<Diagnostic> {
                 ));
 
                 if tag.to_uppercase() == "TODO" {
-                    diagnostics_ref.push(Diagnostic::new(
+                    let invalid_capitalization = Diagnostic::new(
                         InvalidCapitalizationInTodo {
                             tag: String::from(tag),
                         },
                         range,
-                    ));
+                    )
+                    .with_fix(
+                        if should_autofix(autofix, settings, Rule::InvalidCapitalizationInTodo) {
+                            // The TODO regex allows for any number of spaces, so let's find where the first "t"
+                            // or "T" is. We know the unwrap is safe because of the mandatory regex
+                            // match. We'll use position() since "#" is 2 bytes which could throw
+                            // off an implementation that uses byte-indexing.
+                            let first_t_location = comment
+                                .chars()
+                                .position(|c| c.to_string() == "t" || c.to_string() == "T")
+                                .unwrap();
+
+                            Fix::new(vec![Edit::replacement(
+                                "TODO".to_string(),
+                                Location::new(range.location.row(), first_t_location),
+                                Location::new(
+                                    range.location.row(),
+                                    // We know `tag` is 4 bytes long because we've already ensured
+                                    // it's some variant of "TODO"
+                                    first_t_location + tag.len(),
+                                ),
+                            )])
+                        } else {
+                            Fix::empty()
+                        },
+                    );
+
+                    diagnostics_ref.push(invalid_capitalization);
                 }
             }
 
@@ -325,3 +362,7 @@ fn get_regex_error(i: usize, range: &Range, diagnostics: &mut [Diagnostic]) -> O
         _ => None,
     }
 }
+
+fn should_autofix(autofix: flags::Autofix, settings: &Settings, rule: Rule) -> bool {
+    autofix.into() && settings.rules.should_fix(rule)
+}

crates/ruff/src/rules/flake8_todo/snapshots/ruff__rules__flake8_todo__tests__invalid-capitalization-in-todo_TDO006.py.snap

@@ -1,21 +1,38 @@
 ---
 source: crates/ruff/src/rules/flake8_todo/mod.rs
 ---
-TDO006.py:4:1: TDO006 Invalid TODO capitalization: `ToDo` should be `TODO`
+TDO006.py:4:1: TDO006 [*] Invalid TODO capitalization: `ToDo` should be `TODO`
   |
 4 | # TODO(evanrittenhouse): this is a valid TODO
 5 | # TDO006 - error
 6 | # ToDo(evanrittenhouse): invalid capitalization
   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ TDO006
 7 | # todo(evanrittenhouse): another invalid capitalization
   |
+  = help: Fix capitalization in `TODO`
 
-TDO006.py:5:1: TDO006 Invalid TODO capitalization: `todo` should be `TODO`
+ℹ Suggested fix
+1 1 | # TDO006 - accepted
+2 2 | # TODO(evanrittenhouse): this is a valid TODO
+3 3 | # TDO006 - error
+4   |-# ToDo(evanrittenhouse): invalid capitalization
+  4 |+# TODO(evanrittenhouse): invalid capitalization
+5 5 | # todo(evanrittenhouse): another invalid capitalization
+
+TDO006.py:5:1: TDO006 [*] Invalid TODO capitalization: `todo` should be `TODO`
   |
 5 | # TDO006 - error
 6 | # ToDo(evanrittenhouse): invalid capitalization
 7 | # todo(evanrittenhouse): another invalid capitalization
   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ TDO006
   |
+  = help: Fix capitalization in `TODO`
+
+ℹ Suggested fix
+2 2 | # TODO(evanrittenhouse): this is a valid TODO
+3 3 | # TDO006 - error
+4 4 | # ToDo(evanrittenhouse): invalid capitalization
+5   |-# todo(evanrittenhouse): another invalid capitalization
+  5 |+# TODO(evanrittenhouse): another invalid capitalization