Add formatting of type parameters in class and function definitions (#…

…6161)

Part of #5062 
Closes #5931

Implements formatting of a sequence of type parameters in a dedicated
struct for reuse by classes, functions, and type aliases (preparing for
#5929). Adds formatting of type parameters in class and function
definitions — previously, they were just elided.

crates/ruff_python_ast/src/function.rs

@@ -1,6 +1,7 @@
 use crate::node::AnyNodeRef;
 use crate::{
     Decorator, Expr, Identifier, Parameters, Ranged, StmtAsyncFunctionDef, StmtFunctionDef, Suite,
+    TypeParams,
 };
 use ruff_text_size::TextRange;
 
@@ -79,6 +80,13 @@ impl<'a> AnyFunctionDefinition<'a> {
         }
     }
 
+    pub fn type_params(self) -> Option<&'a TypeParams> {
+        match self {
+            Self::FunctionDefinition(definition) => definition.type_params.as_ref(),
+            Self::AsyncFunctionDefinition(definition) => definition.type_params.as_ref(),
+        }
+    }
+
     /// Returns `true` if this is [`Self::AsyncFunctionDefinition`]
     pub const fn is_async(self) -> bool {
         matches!(self, Self::AsyncFunctionDefinition(_))

crates/ruff_python_ast/src/node.rs

@@ -2880,6 +2880,34 @@ impl AstNode for Decorator {
         AnyNode::from(self)
     }
 }
+impl AstNode for ast::TypeParams {
+    fn cast(kind: AnyNode) -> Option<Self>
+    where
+        Self: Sized,
+    {
+        if let AnyNode::TypeParams(node) = kind {
+            Some(node)
+        } else {
+            None
+        }
+    }
+
+    fn cast_ref(kind: AnyNodeRef) -> Option<&Self> {
+        if let AnyNodeRef::TypeParams(node) = kind {
+            Some(node)
+        } else {
+            None
+        }
+    }
+
+    fn as_any_node_ref(&self) -> AnyNodeRef {
+        AnyNodeRef::from(self)
+    }
+
+    fn into_any_node(self) -> AnyNode {
+        AnyNode::from(self)
+    }
+}
 impl AstNode for ast::TypeParamTypeVar {
     fn cast(kind: AnyNode) -> Option<Self>
     where
@@ -3531,6 +3559,11 @@ impl From<Decorator> for AnyNode {
         AnyNode::Decorator(node)
     }
 }
+impl From<TypeParams> for AnyNode {
+    fn from(node: TypeParams) -> Self {
+        AnyNode::TypeParams(node)
+    }
+}
 impl From<TypeParamTypeVar> for AnyNode {
     fn from(node: TypeParamTypeVar) -> Self {
         AnyNode::TypeParamTypeVar(node)
@@ -4804,6 +4837,11 @@ impl<'a> From<&'a Decorator> for AnyNodeRef<'a> {
     }
 }
 
+impl<'a> From<&'a ast::TypeParams> for AnyNodeRef<'a> {
+    fn from(node: &'a ast::TypeParams) -> Self {
+        AnyNodeRef::TypeParams(node)
+    }
+}
 impl<'a> From<&'a TypeParamTypeVar> for AnyNodeRef<'a> {
     fn from(node: &'a TypeParamTypeVar) -> Self {
         AnyNodeRef::TypeParamTypeVar(node)

crates/ruff_python_formatter/generate.py

@@ -54,6 +54,7 @@ def rustfmt(code: str) -> str:
     "expr": "expression",
     "stmt": "statement",
     "pattern": "pattern",
+    "type_param": "type_param",
     "other": "other",
 }
 

crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/list.py

@@ -21,3 +21,14 @@
     aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa,
     aksjdhflsakhdflkjsadlfajkslhfdkjsaldajlahflashdfljahlfksajlhfajfjfsaahflakjslhdfkjalhdskjfa
 ]
+
+# Comment placement in non-empty lists
+c1 = [ # trailing open bracket
+    # leading item
+    1,
+
+    # between
+
+    2, # trailing item
+    # leading close bracket
+] # trailing close bracket

crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/class_definition.py

@@ -43,6 +43,14 @@ class TestTrailingComment2: # trailing comment
     pass
 
 
+class TestTrailingComment3[T]: # trailing comment
+    pass
+
+
+class TestTrailingComment4[T](A): # trailing comment
+    pass
+
+
 class Test:
     """Docstring"""
 
@@ -144,3 +152,57 @@ class AltCLIPOutput(
     # Copied from transformers.models.clip.modeling_clip.CLIPOutput with CLIP->AltCLIP
 ):
     ...
+
+
+class TestTypeParams[Aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa, Bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb, Cccccccccccccccccccccc]:
+    pass
+
+
+class TestTypeParams[Aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa, *Bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb, **Cccccccccccccccccccccc]:
+    pass
+
+
+class TestTypeParams[Aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa]:
+    pass
+
+
+class TestTypeParams[*Aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa]:
+    pass
+
+
+class TestTypeParams[**Aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa]:
+    pass
+
+
+class TestTypeParams[**P, *Ts, T]:
+    pass
+
+
+class TestTypeParams[ # trailing bracket comment
+    # leading comment
+    A,
+
+    # in between comment
+
+    B,
+    # another leading comment
+    C,
+    D, # trailing comment
+    # leading bracket comment
+]:
+    pass
+
+
+
+class TestTypeParams[Aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa](B, C, D):
+    pass
+
+
+
+class TestTypeParams[Aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa](Bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb, Cccccccccccccccccccccccc, Ddddddddddddddddddddddddd):
+    pass
+
+
+
+class TestTypeParams[A, B, C](meta=Aaaaaaaaaaaaaaaaaaaaaa):
+    pass

crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/function.py

@@ -66,6 +66,52 @@ def argument_with_long_type_annotation(
 
 def test(): ...
 
+# Type parameter empty line spacing
+def test[
+    # comment
+    A,
+
+    # another
+
+    B,
+](): ...
+
+
+# Type parameter comments
+def type_param_comments[ # trailing bracket comment
+    # leading comment
+    A,
+
+    # in between comment
+
+    B,
+    # another leading comment
+    C,
+    D, # trailing comment
+    # leading bracket comment
+]():
+    # body comment
+    pass
+
+
+# Note empty type parameters is not valid syntax, e.g.
+# def test[](): ...
+
+
+# Different type parameter wrappings
+
+def single_line[Aaaaaaaaaaaaaaaaaaaaaaaaaaaaa, Bbbbbbbbbbbbbbb, Ccccccccccccccccc]():
+    pass
+
+def params_on_their_own_line[Aaaaaaaaaaaaaaaaaaaaaaaaaaaaa, Bbbbbbbbbbbbbbb, Ccccccccccc, Ddddddddddddd, Eeeeeeee]():
+    pass
+
+def param_per_line[Aaaaaaaaaaaaaaaaaaaaaaaaaaaaa, Bbbbbbbbbbbbbbb, Ccccccccccccccccc, Ddddddddddddd, Eeeeeeeeeeeeeeeee, ffffffffffff]():
+    pass
+
+def single_line_trailing_comma[A, B, C,]():
+    pass
+
 # Comment
 def with_leading_comment(): ...
 

crates/ruff_python_formatter/src/comments/placement.rs

@@ -4,7 +4,7 @@ use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::whitespace::indentation;
 use ruff_python_ast::{
     self as ast, Arguments, Comprehension, Expr, ExprAttribute, ExprBinOp, ExprIfExp, ExprSlice,
-    ExprStarred, MatchCase, Parameters, Ranged,
+    ExprStarred, MatchCase, Parameters, Ranged, TypeParams,
 };
 use ruff_python_trivia::{
     indentation_at_offset, PythonWhitespace, SimpleToken, SimpleTokenKind, SimpleTokenizer,
@@ -85,6 +85,7 @@ pub(super) fn place_comment<'a>(
             handle_leading_class_with_decorators_comment(comment, class_def)
         }
         AnyNodeRef::StmtImportFrom(import_from) => handle_import_from_comment(comment, import_from),
+        AnyNodeRef::TypeParams(type_params) => handle_type_params_comment(comment, type_params),
         _ => CommentPlacement::Default(comment),
     }
 }
@@ -563,6 +564,51 @@ fn handle_own_line_comment_after_branch<'a>(
     }
 }
 
+/// Attach an enclosed end-of-line comment to a set of [`TypeParams`].
+///
+/// For example, given:
+/// ```python
+/// type foo[  # comment
+///    bar,
+/// ] = ...
+/// ```
+///
+/// The comment will be attached to the [`TypeParams`] node as a dangling comment, to ensure
+/// that it remains on the same line as open bracket.
+fn handle_type_params_comment<'a>(
+    comment: DecoratedComment<'a>,
+    type_params: &'a TypeParams,
+) -> CommentPlacement<'a> {
+    // The comment needs to be on the same line, but before the first type param. For example, we want
+    // to treat this as a dangling comment:
+    // ```python
+    // type foo[  # comment
+    //     bar,
+    //     baz,
+    //     qux,
+    // ]
+    // ```
+    // However, this should _not_ be treated as a dangling comment:
+    // ```python
+    // type foo[bar,  # comment
+    //     baz,
+    //     qux,
+    // ] = ...
+    // ```
+    // Thus, we check whether the comment is an end-of-line comment _between_ the start of the
+    // statement and the first type param. If so, the only possible position is immediately following
+    // the open parenthesis.
+    if comment.line_position().is_end_of_line() {
+        if let Some(first_type_param) = type_params.type_params.first() {
+            if type_params.start() < comment.start() && comment.end() < first_type_param.start() {
+                return CommentPlacement::dangling(comment.enclosing_node(), comment);
+            }
+        }
+    }
+
+    CommentPlacement::Default(comment)
+}
+
 /// Attaches comments for the positional-only parameters separator `/` or the keywords-only
 /// parameters separator `*` as dangling comments to the enclosing [`Parameters`] node.
 ///

crates/ruff_python_formatter/src/comments/visitor.rs

@@ -3,7 +3,7 @@ use std::iter::Peekable;
 use ruff_python_ast::{
     Alias, Arguments, Comprehension, Decorator, ElifElseClause, ExceptHandler, Expr, Keyword,
     MatchCase, Mod, Parameter, ParameterWithDefault, Parameters, Pattern, Ranged, Stmt, TypeParam,
-    WithItem,
+    TypeParams, WithItem,
 };
 use ruff_text_size::{TextRange, TextSize};
 
@@ -301,6 +301,13 @@ impl<'ast> PreorderVisitor<'ast> for CommentsVisitor<'ast> {
         self.finish_node(elif_else_clause);
     }
 
+    fn visit_type_params(&mut self, type_params: &'ast TypeParams) {
+        if self.start_node(type_params).is_traverse() {
+            walk_type_params(self, type_params);
+        }
+        self.finish_node(type_params);
+    }
+
     fn visit_type_param(&mut self, type_param: &'ast TypeParam) {
         if self.start_node(type_param).is_traverse() {
             walk_type_param(self, type_param);