Optimise, and fix bug regarding indented class/function definitions

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

@@ -114,13 +114,17 @@ def pep_696_bad[A:object="foo"[::-1], B:object =[[["foo", "bar"]]], C:object= by
     pass
 
 class PEP696Bad[A:object="foo"[::-1], B:object =[[["foo", "bar"]]], C:object= bytes]:
-    pass
+    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]():
-    pass
-
-class PEP696BadWithNonEmptyBases[A:object="foo"[::-1], B:object =[[["foo", "bar"]]], C:object= bytes](object, something_dynamic[x::-1]):
-    pass
+    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](

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

@@ -73,4 +73,5 @@ def pep_696_good[A = int, B: object = str, C:object = memoryview]():
     pass
 
 class PEP696Good[A = int, B: object = str, C:object = memoryview]:
-    pass
+    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, TypeParamsState};
+use super::{DefinitionState, LogicalLine};
 
 /// ## What it does
 /// Checks for missing whitespace after `,`, `;`, and `:`.
@@ -42,13 +42,13 @@ impl AlwaysFixableViolation for MissingWhitespace {
 /// E231
 pub(crate) fn missing_whitespace(line: &LogicalLine, context: &mut LogicalLinesContext) {
     let mut fstrings = 0u32;
-    let mut type_params_state = TypeParamsState::new();
+    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();
-        type_params_state.visit_token_kind(kind);
+        definition_state.visit_token_kind(kind);
         match kind {
             TokenKind::FStringStart => fstrings += 1,
             TokenKind::FStringEnd => fstrings = fstrings.saturating_sub(1),
@@ -88,7 +88,7 @@ pub(crate) fn missing_whitespace(line: &LogicalLine, context: &mut LogicalLinesC
                         match (kind, next_token.kind()) {
                             (TokenKind::Colon, _)
                                 if matches!(brackets.last(), Some(TokenKind::Lsqb))
-                                    && !(type_params_state.in_type_params()
+                                    && !(definition_state.in_type_params()
                                         && brackets.len() == 1) =>
                             {
                                 continue; // Slice syntax, no space required

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

@@ -470,50 +470,72 @@ struct Line {
     tokens_end: u32,
 }
 
-/// Keeps track of whether we are currently visiting the [type parameters]
-/// of a class or function definition in a [`LogicalLine`].
+/// 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 [`TypeParamsState::visit_token_kind`] on the [`TokenKind`] of each
+/// 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 TypeParamsState {
-    BeforeClassOrDefKeyword,
-    BeforeTypeParams,
-    InTypeParams { inner_square_brackets: u32 },
-    TypeParamsEnded,
+enum DefinitionState {
+    InClass(TypeParamsState),
+    InFunction(TypeParamsState),
+    NotInClassOrFunction,
 }
 
-impl TypeParamsState {
-    const fn new() -> Self {
-        Self::BeforeClassOrDefKeyword
+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_class_or_function_def(self) -> bool {
-        !matches!(self, Self::BeforeClassOrDefKeyword)
+    const fn in_function_definition(self) -> bool {
+        matches!(self, Self::InFunction(_))
     }
 
-    const fn before_type_params(self) -> bool {
-        matches!(self, Self::BeforeTypeParams)
+    const fn type_params_state(self) -> Option<TypeParamsState> {
+        match self {
+            Self::InClass(state) | Self::InFunction(state) => Some(state),
+            Self::NotInClassOrFunction => None,
+        }
     }
 
-    const fn in_type_params(self) -> bool {
-        matches!(self, Self::InTypeParams { .. })
+    fn in_type_params(self) -> bool {
+        matches!(
+            self.type_params_state(),
+            Some(TypeParamsState::InTypeParams { .. })
+        )
     }
 
-    fn visit_token_kind(&mut self, token: TokenKind) {
-        match token {
-            TokenKind::Class | TokenKind::Def if !self.in_class_or_function_def() => {
-                *self = TypeParamsState::BeforeTypeParams;
+    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
             }
-            TokenKind::Lpar if self.before_type_params() => {
-                *self = TypeParamsState::TypeParamsEnded;
+            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 self {
-                TypeParamsState::BeforeClassOrDefKeyword | TypeParamsState::TypeParamsEnded => {}
+            TokenKind::Lsqb => match type_params_state_mut {
+                TypeParamsState::TypeParamsEnded => {}
                 TypeParamsState::BeforeTypeParams => {
-                    *self = TypeParamsState::InTypeParams {
+                    *type_params_state_mut = TypeParamsState::InTypeParams {
                         inner_square_brackets: 0,
                     };
                 }
@@ -524,10 +546,10 @@ impl TypeParamsState {
             TokenKind::Rsqb => {
                 if let TypeParamsState::InTypeParams {
                     inner_square_brackets,
-                } = self
+                } = type_params_state_mut
                 {
                     if *inner_square_brackets == 0 {
-                        *self = TypeParamsState::TypeParamsEnded;
+                        *type_params_state_mut = TypeParamsState::TypeParamsEnded;
                     } else {
                         *inner_square_brackets -= 1;
                     }
@@ -538,6 +560,22 @@ impl TypeParamsState {
     }
 }
 
+#[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,9 +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, TypeParamsState,
-};
+use crate::rules::pycodestyle::rules::logical_lines::{DefinitionState, LogicalLine};
 
 /// ## What it does
 /// Checks for missing whitespace around the equals sign in an unannotated
@@ -86,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,
@@ -108,13 +94,12 @@ pub(crate) fn whitespace_around_named_parameter_equals(
     let mut annotated_func_arg = false;
     let mut prev_end = TextSize::default();
 
-    let mut type_params_state = TypeParamsState::new();
-    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 token_kind = token.kind();
-        type_params_state.visit_token_kind(token_kind);
+        definition_state.visit_token_kind(token_kind);
         match token_kind {
             TokenKind::NonLogicalNewline => continue,
             TokenKind::FStringStart => fstrings += 1,
@@ -128,16 +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 type_params_state.in_type_params() || (parens > 0 && fstrings == 0) =>
+                if definition_state.in_type_params() || (parens > 0 && fstrings == 0) =>
             {
-                if type_params_state.in_type_params() || (annotated_func_arg && parens == 1) {
+                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 =