[pycodestyle] Fix whitespace-related false positives and false nega…

…tives inside type-parameter lists (#13704)

crates/ruff_linter/resources/test/fixtures/pycodestyle/E23.py

@@ -104,3 +104,41 @@ def main() -> None:
         ]
     }
 ]
+
+# Should be E231 errors on all of these type parameters and function parameters, but not on their (strange) defaults
+def pep_696_bad[A:object="foo"[::-1], B:object =[[["foo", "bar"]]], C:object= bytes](
+    x:A = "foo"[::-1],
+    y:B = [[["foo", "bar"]]],
+    z:object = "fooo",
+):
+    pass
+
+class PEP696Bad[A:object="foo"[::-1], B:object =[[["foo", "bar"]]], C:object= bytes]:
+    def pep_696_bad_method[A:object="foo"[::-1], B:object =[[["foo", "bar"]]], C:object= bytes](
+        self,
+        x:A = "foo"[::-1],
+        y:B = [[["foo", "bar"]]],
+        z:object = "fooo",
+    ):
+        pass
+
+class PEP696BadWithEmptyBases[A:object="foo"[::-1], B:object =[[["foo", "bar"]]], C:object= bytes]():
+    class IndentedPEP696BadWithNonEmptyBases[A:object="foo"[::-1], B:object =[[["foo", "bar"]]], C:object= bytes](object, something_dynamic[x::-1]):
+        pass
+
+# Should be no E231 errors on any of these:
+def pep_696_good[A: object="foo"[::-1], B: object =[[["foo", "bar"]]], C: object= bytes](
+    x: A = "foo"[::-1],
+    y: B = [[["foo", "bar"]]],
+    z: object = "fooo",
+):
+    pass
+
+class PEP696Good[A: object="foo"[::-1], B: object =[[["foo", "bar"]]], C: object= bytes]:
+    pass
+
+class PEP696GoodWithEmptyBases[A: object="foo"[::-1], B: object =[[["foo", "bar"]]], C: object= bytes]():
+    pass
+
+class PEP696GoodWithNonEmptyBases[A: object="foo"[::-1], B: object =[[["foo", "bar"]]], C: object= bytes](object, something_dynamic[x::-1]):
+    pass

crates/ruff_linter/resources/test/fixtures/pycodestyle/E25.py

@@ -59,3 +59,19 @@ def add(a: int = _default(name='f')):
 print(f"{foo = }")
 # ...but then it creates false negatives for now
 print(f"{foo(a = 1)}")
+
+# There should be at least one E251 diagnostic for each type parameter here:
+def pep_696_bad[A=int, B =str, C= bool, D:object=int, E: object=str, F: object =bool, G: object= bytes]():
+    pass
+
+class PEP696Bad[A=int, B =str, C= bool, D:object=int, E: object=str, F: object =bool, G: object= bytes]:
+    pass
+
+# The last of these should cause us to emit E231,
+# but E231 isn't tested by this fixture:
+def pep_696_good[A = int, B: object = str, C:object = memoryview]():
+    pass
+
+class PEP696Good[A = int, B: object = str, C:object = memoryview]:
+    def pep_696_good_method[A = int, B: object = str, C:object = memoryview](self):
+        pass

crates/ruff_linter/src/rules/pycodestyle/rules/logical_lines/missing_whitespace.rs

@@ -6,7 +6,7 @@ use ruff_text_size::Ranged;
 
 use crate::checkers::logical_lines::LogicalLinesContext;
 
-use super::LogicalLine;
+use super::{DefinitionState, LogicalLine};
 
 /// ## What it does
 /// Checks for missing whitespace after `,`, `;`, and `:`.
@@ -28,22 +28,10 @@ pub struct MissingWhitespace {
     token: TokenKind,
 }
 
-impl MissingWhitespace {
-    fn token_text(&self) -> char {
-        match self.token {
-            TokenKind::Colon => ':',
-            TokenKind::Semi => ';',
-            TokenKind::Comma => ',',
-            _ => unreachable!(),
-        }
-    }
-}
-
 impl AlwaysFixableViolation for MissingWhitespace {
     #[derive_message_formats]
     fn message(&self) -> String {
-        let token = self.token_text();
-        format!("Missing whitespace after '{token}'")
+        format!("Missing whitespace after {}", self.token)
     }
 
     fn fix_title(&self) -> String {
@@ -54,11 +42,13 @@ impl AlwaysFixableViolation for MissingWhitespace {
 /// E231
 pub(crate) fn missing_whitespace(line: &LogicalLine, context: &mut LogicalLinesContext) {
     let mut fstrings = 0u32;
+    let mut definition_state = DefinitionState::from_tokens(line.tokens());
     let mut brackets = Vec::new();
     let mut iter = line.tokens().iter().peekable();
 
     while let Some(token) = iter.next() {
         let kind = token.kind();
+        definition_state.visit_token_kind(kind);
         match kind {
             TokenKind::FStringStart => fstrings += 1,
             TokenKind::FStringEnd => fstrings = fstrings.saturating_sub(1),
@@ -97,7 +87,9 @@ pub(crate) fn missing_whitespace(line: &LogicalLine, context: &mut LogicalLinesC
                     if let Some(next_token) = iter.peek() {
                         match (kind, next_token.kind()) {
                             (TokenKind::Colon, _)
-                                if matches!(brackets.last(), Some(TokenKind::Lsqb)) =>
+                                if matches!(brackets.last(), Some(TokenKind::Lsqb))
+                                    && !(definition_state.in_type_params()
+                                        && brackets.len() == 1) =>
                             {
                                 continue; // Slice syntax, no space required
                             }
@@ -111,13 +103,10 @@ pub(crate) fn missing_whitespace(line: &LogicalLine, context: &mut LogicalLinesC
                         }
                     }
 
-                    let mut diagnostic =
+                    let diagnostic =
                         Diagnostic::new(MissingWhitespace { token: kind }, token.range());
-                    diagnostic.set_fix(Fix::safe_edit(Edit::insertion(
-                        " ".to_string(),
-                        token.end(),
-                    )));
-                    context.push_diagnostic(diagnostic);
+                    let fix = Fix::safe_edit(Edit::insertion(" ".to_string(), token.end()));
+                    context.push_diagnostic(diagnostic.with_fix(fix));
                 }
             }
             _ => {}

crates/ruff_linter/src/rules/pycodestyle/rules/logical_lines/mod.rs

@@ -470,6 +470,112 @@ struct Line {
     tokens_end: u32,
 }
 
+/// Keeps track of whether we are currently visiting a class or function definition in a
+/// [`LogicalLine`]. If we are visiting a class or function, the enum also keeps track
+/// of the [type parameters] of the class/function.
+///
+/// Call [`DefinitionState::visit_token_kind`] on the [`TokenKind`] of each
+/// successive [`LogicalLineToken`] to ensure the state remains up to date.
+///
+/// [type parameters]: https://docs.python.org/3/reference/compound_stmts.html#type-params
+#[derive(Debug, Clone, Copy)]
+enum DefinitionState {
+    InClass(TypeParamsState),
+    InFunction(TypeParamsState),
+    NotInClassOrFunction,
+}
+
+impl DefinitionState {
+    fn from_tokens<'a>(tokens: impl IntoIterator<Item = &'a LogicalLineToken>) -> Self {
+        let mut token_kinds = tokens.into_iter().map(LogicalLineToken::kind);
+        while let Some(token_kind) = token_kinds.next() {
+            let state = match token_kind {
+                TokenKind::Indent | TokenKind::Dedent => continue,
+                TokenKind::Class => Self::InClass(TypeParamsState::default()),
+                TokenKind::Def => Self::InFunction(TypeParamsState::default()),
+                TokenKind::Async if matches!(token_kinds.next(), Some(TokenKind::Def)) => {
+                    Self::InFunction(TypeParamsState::default())
+                }
+                _ => Self::NotInClassOrFunction,
+            };
+            return state;
+        }
+        Self::NotInClassOrFunction
+    }
+
+    const fn in_function_definition(self) -> bool {
+        matches!(self, Self::InFunction(_))
+    }
+
+    const fn type_params_state(self) -> Option<TypeParamsState> {
+        match self {
+            Self::InClass(state) | Self::InFunction(state) => Some(state),
+            Self::NotInClassOrFunction => None,
+        }
+    }
+
+    fn in_type_params(self) -> bool {
+        matches!(
+            self.type_params_state(),
+            Some(TypeParamsState::InTypeParams { .. })
+        )
+    }
+
+    fn visit_token_kind(&mut self, token_kind: TokenKind) {
+        let type_params_state_mut = match self {
+            Self::InClass(type_params_state) | Self::InFunction(type_params_state) => {
+                type_params_state
+            }
+            Self::NotInClassOrFunction => return,
+        };
+        match token_kind {
+            TokenKind::Lpar if type_params_state_mut.before_type_params() => {
+                *type_params_state_mut = TypeParamsState::TypeParamsEnded;
+            }
+            TokenKind::Lsqb => match type_params_state_mut {
+                TypeParamsState::TypeParamsEnded => {}
+                TypeParamsState::BeforeTypeParams => {
+                    *type_params_state_mut = TypeParamsState::InTypeParams {
+                        inner_square_brackets: 0,
+                    };
+                }
+                TypeParamsState::InTypeParams {
+                    inner_square_brackets,
+                } => *inner_square_brackets += 1,
+            },
+            TokenKind::Rsqb => {
+                if let TypeParamsState::InTypeParams {
+                    inner_square_brackets,
+                } = type_params_state_mut
+                {
+                    if *inner_square_brackets == 0 {
+                        *type_params_state_mut = TypeParamsState::TypeParamsEnded;
+                    } else {
+                        *inner_square_brackets -= 1;
+                    }
+                }
+            }
+            _ => {}
+        }
+    }
+}
+
+#[derive(Debug, Clone, Copy, Default)]
+enum TypeParamsState {
+    #[default]
+    BeforeTypeParams,
+    InTypeParams {
+        inner_square_brackets: u32,
+    },
+    TypeParamsEnded,
+}
+
+impl TypeParamsState {
+    const fn before_type_params(self) -> bool {
+        matches!(self, Self::BeforeTypeParams)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use ruff_python_parser::parse_module;

crates/ruff_linter/src/rules/pycodestyle/rules/logical_lines/whitespace_around_named_parameter_equals.rs

@@ -4,7 +4,7 @@ use ruff_python_parser::TokenKind;
 use ruff_text_size::{Ranged, TextRange, TextSize};
 
 use crate::checkers::logical_lines::LogicalLinesContext;
-use crate::rules::pycodestyle::rules::logical_lines::{LogicalLine, LogicalLineToken};
+use crate::rules::pycodestyle::rules::logical_lines::{DefinitionState, LogicalLine};
 
 /// ## What it does
 /// Checks for missing whitespace around the equals sign in an unannotated
@@ -84,18 +84,6 @@ impl AlwaysFixableViolation for MissingWhitespaceAroundParameterEquals {
     }
 }
 
-fn is_in_def(tokens: &[LogicalLineToken]) -> bool {
-    for token in tokens {
-        match token.kind() {
-            TokenKind::Async | TokenKind::Indent | TokenKind::Dedent => continue,
-            TokenKind::Def => return true,
-            _ => return false,
-        }
-    }
-
-    false
-}
-
 /// E251, E252
 pub(crate) fn whitespace_around_named_parameter_equals(
     line: &LogicalLine,
@@ -106,17 +94,14 @@ pub(crate) fn whitespace_around_named_parameter_equals(
     let mut annotated_func_arg = false;
     let mut prev_end = TextSize::default();
 
-    let in_def = is_in_def(line.tokens());
+    let mut definition_state = DefinitionState::from_tokens(line.tokens());
     let mut iter = line.tokens().iter().peekable();
 
     while let Some(token) = iter.next() {
         let kind = token.kind();
-
-        if kind == TokenKind::NonLogicalNewline {
-            continue;
-        }
-
+        definition_state.visit_token_kind(kind);
         match kind {
+            TokenKind::NonLogicalNewline => continue,
             TokenKind::FStringStart => fstrings += 1,
             TokenKind::FStringEnd => fstrings = fstrings.saturating_sub(1),
             TokenKind::Lpar | TokenKind::Lsqb => {
@@ -128,15 +113,16 @@ pub(crate) fn whitespace_around_named_parameter_equals(
                     annotated_func_arg = false;
                 }
             }
-
-            TokenKind::Colon if parens == 1 && in_def => {
+            TokenKind::Colon if parens == 1 && definition_state.in_function_definition() => {
                 annotated_func_arg = true;
             }
             TokenKind::Comma if parens == 1 => {
                 annotated_func_arg = false;
             }
-            TokenKind::Equal if parens > 0 && fstrings == 0 => {
-                if annotated_func_arg && parens == 1 {
+            TokenKind::Equal
+                if definition_state.in_type_params() || (parens > 0 && fstrings == 0) =>
+            {
+                if definition_state.in_type_params() || (annotated_func_arg && parens == 1) {
                     let start = token.start();
                     if start == prev_end && prev_end != TextSize::new(0) {
                         let mut diagnostic =