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

Part of #5062 
Closes https://github.com/astral-sh/ruff/issues/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_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);

crates/ruff_python_formatter/src/generated.rs

@@ -2939,3 +2939,163 @@ impl<'ast> IntoFormat<PyFormatContext<'ast>> for ast::ElifElseClause {
         )
     }
 }
+
+impl FormatRule<ast::TypeParams, PyFormatContext<'_>>
+    for crate::type_param::type_params::FormatTypeParams
+{
+    #[inline]
+    fn fmt(
+        &self,
+        node: &ast::TypeParams,
+        f: &mut Formatter<PyFormatContext<'_>>,
+    ) -> FormatResult<()> {
+        FormatNodeRule::<ast::TypeParams>::fmt(self, node, f)
+    }
+}
+impl<'ast> AsFormat<PyFormatContext<'ast>> for ast::TypeParams {
+    type Format<'a> = FormatRefWithRule<
+        'a,
+        ast::TypeParams,
+        crate::type_param::type_params::FormatTypeParams,
+        PyFormatContext<'ast>,
+    >;
+    fn format(&self) -> Self::Format<'_> {
+        FormatRefWithRule::new(
+            self,
+            crate::type_param::type_params::FormatTypeParams::default(),
+        )
+    }
+}
+impl<'ast> IntoFormat<PyFormatContext<'ast>> for ast::TypeParams {
+    type Format = FormatOwnedWithRule<
+        ast::TypeParams,
+        crate::type_param::type_params::FormatTypeParams,
+        PyFormatContext<'ast>,
+    >;
+    fn into_format(self) -> Self::Format {
+        FormatOwnedWithRule::new(
+            self,
+            crate::type_param::type_params::FormatTypeParams::default(),
+        )
+    }
+}
+
+impl FormatRule<ast::TypeParamTypeVar, PyFormatContext<'_>>
+    for crate::type_param::type_param_type_var::FormatTypeParamTypeVar
+{
+    #[inline]
+    fn fmt(
+        &self,
+        node: &ast::TypeParamTypeVar,
+        f: &mut Formatter<PyFormatContext<'_>>,
+    ) -> FormatResult<()> {
+        FormatNodeRule::<ast::TypeParamTypeVar>::fmt(self, node, f)
+    }
+}
+impl<'ast> AsFormat<PyFormatContext<'ast>> for ast::TypeParamTypeVar {
+    type Format<'a> = FormatRefWithRule<
+        'a,
+        ast::TypeParamTypeVar,
+        crate::type_param::type_param_type_var::FormatTypeParamTypeVar,
+        PyFormatContext<'ast>,
+    >;
+    fn format(&self) -> Self::Format<'_> {
+        FormatRefWithRule::new(
+            self,
+            crate::type_param::type_param_type_var::FormatTypeParamTypeVar::default(),
+        )
+    }
+}
+impl<'ast> IntoFormat<PyFormatContext<'ast>> for ast::TypeParamTypeVar {
+    type Format = FormatOwnedWithRule<
+        ast::TypeParamTypeVar,
+        crate::type_param::type_param_type_var::FormatTypeParamTypeVar,
+        PyFormatContext<'ast>,
+    >;
+    fn into_format(self) -> Self::Format {
+        FormatOwnedWithRule::new(
+            self,
+            crate::type_param::type_param_type_var::FormatTypeParamTypeVar::default(),
+        )
+    }
+}
+
+impl FormatRule<ast::TypeParamTypeVarTuple, PyFormatContext<'_>>
+    for crate::type_param::type_param_type_var_tuple::FormatTypeParamTypeVarTuple
+{
+    #[inline]
+    fn fmt(
+        &self,
+        node: &ast::TypeParamTypeVarTuple,
+        f: &mut Formatter<PyFormatContext<'_>>,
+    ) -> FormatResult<()> {
+        FormatNodeRule::<ast::TypeParamTypeVarTuple>::fmt(self, node, f)
+    }
+}
+impl<'ast> AsFormat<PyFormatContext<'ast>> for ast::TypeParamTypeVarTuple {
+    type Format<'a> = FormatRefWithRule<
+        'a,
+        ast::TypeParamTypeVarTuple,
+        crate::type_param::type_param_type_var_tuple::FormatTypeParamTypeVarTuple,
+        PyFormatContext<'ast>,
+    >;
+    fn format(&self) -> Self::Format<'_> {
+        FormatRefWithRule::new(
+            self,
+            crate::type_param::type_param_type_var_tuple::FormatTypeParamTypeVarTuple::default(),
+        )
+    }
+}
+impl<'ast> IntoFormat<PyFormatContext<'ast>> for ast::TypeParamTypeVarTuple {
+    type Format = FormatOwnedWithRule<
+        ast::TypeParamTypeVarTuple,
+        crate::type_param::type_param_type_var_tuple::FormatTypeParamTypeVarTuple,
+        PyFormatContext<'ast>,
+    >;
+    fn into_format(self) -> Self::Format {
+        FormatOwnedWithRule::new(
+            self,
+            crate::type_param::type_param_type_var_tuple::FormatTypeParamTypeVarTuple::default(),
+        )
+    }
+}
+
+impl FormatRule<ast::TypeParamParamSpec, PyFormatContext<'_>>
+    for crate::type_param::type_param_param_spec::FormatTypeParamParamSpec
+{
+    #[inline]
+    fn fmt(
+        &self,
+        node: &ast::TypeParamParamSpec,
+        f: &mut Formatter<PyFormatContext<'_>>,
+    ) -> FormatResult<()> {
+        FormatNodeRule::<ast::TypeParamParamSpec>::fmt(self, node, f)
+    }
+}
+impl<'ast> AsFormat<PyFormatContext<'ast>> for ast::TypeParamParamSpec {
+    type Format<'a> = FormatRefWithRule<
+        'a,
+        ast::TypeParamParamSpec,
+        crate::type_param::type_param_param_spec::FormatTypeParamParamSpec,
+        PyFormatContext<'ast>,
+    >;
+    fn format(&self) -> Self::Format<'_> {
+        FormatRefWithRule::new(
+            self,
+            crate::type_param::type_param_param_spec::FormatTypeParamParamSpec::default(),
+        )
+    }
+}
+impl<'ast> IntoFormat<PyFormatContext<'ast>> for ast::TypeParamParamSpec {
+    type Format = FormatOwnedWithRule<
+        ast::TypeParamParamSpec,
+        crate::type_param::type_param_param_spec::FormatTypeParamParamSpec,
+        PyFormatContext<'ast>,
+    >;
+    fn into_format(self) -> Self::Format {
+        FormatOwnedWithRule::new(
+            self,
+            crate::type_param::type_param_param_spec::FormatTypeParamParamSpec::default(),
+        )
+    }
+}

crates/ruff_python_formatter/src/lib.rs

@@ -34,6 +34,7 @@ pub(crate) mod other;
 pub(crate) mod pattern;
 mod prelude;
 pub(crate) mod statement;
+pub(crate) mod type_param;
 
 include!("../../ruff_formatter/shared_traits.rs");
 

crates/ruff_python_formatter/src/statement/stmt_class_def.rs

@@ -1,10 +1,11 @@
-use ruff_formatter::write;
+use ruff_formatter::{write, Buffer};
 use ruff_python_ast::{Ranged, StmtClassDef};
 use ruff_python_trivia::{lines_after, skip_trailing_trivia};
 
 use crate::comments::{leading_comments, trailing_comments};
 use crate::prelude::*;
 use crate::statement::suite::SuiteKind;
+use crate::FormatNodeRule;
 
 #[derive(Default)]
 pub struct FormatStmtClassDef;
@@ -16,7 +17,7 @@ impl FormatNodeRule<StmtClassDef> for FormatStmtClassDef {
             name,
             arguments,
             body,
-            type_params: _,
+            type_params,
             decorator_list,
         } = item;
 
@@ -58,6 +59,10 @@ impl FormatNodeRule<StmtClassDef> for FormatStmtClassDef {
 
         write!(f, [text("class"), space(), name.format()])?;
 
+        if let Some(type_params) = type_params.as_deref() {
+            write!(f, [type_params.format()])?;
+        }
+
         if let Some(arguments) = arguments.as_deref() {
             // Drop empty parentheses, e.g., in:
             // ```python

crates/ruff_python_formatter/src/statement/stmt_function_def.rs

@@ -80,15 +80,13 @@ impl FormatRule<AnyFunctionDefinition<'_>, PyFormatContext<'_>> for FormatAnyFun
 
         let name = item.name();
 
-        write!(
-            f,
-            [
-                text("def"),
-                space(),
-                name.format(),
-                item.arguments().format(),
-            ]
-        )?;
+        write!(f, [text("def"), space(), name.format()])?;
+
+        if let Some(type_params) = item.type_params() {
+            write!(f, [type_params.format()])?;
+        }
+
+        write!(f, [item.arguments().format()])?;
 
         if let Some(return_annotation) = item.returns() {
             write!(

crates/ruff_python_formatter/src/type_param/mod.rs

@@ -0,0 +1,37 @@
+use crate::prelude::*;
+use ruff_formatter::{FormatOwnedWithRule, FormatRefWithRule};
+use ruff_python_ast::TypeParam;
+
+pub(crate) mod type_param_param_spec;
+pub(crate) mod type_param_type_var;
+pub(crate) mod type_param_type_var_tuple;
+pub(crate) mod type_params;
+
+#[derive(Default)]
+pub struct FormatTypeParam;
+
+impl FormatRule<TypeParam, PyFormatContext<'_>> for FormatTypeParam {
+    fn fmt(&self, item: &TypeParam, f: &mut PyFormatter) -> FormatResult<()> {
+        match item {
+            TypeParam::TypeVar(x) => x.format().fmt(f),
+            TypeParam::TypeVarTuple(x) => x.format().fmt(f),
+            TypeParam::ParamSpec(x) => x.format().fmt(f),
+        }
+    }
+}
+
+impl<'ast> AsFormat<PyFormatContext<'ast>> for TypeParam {
+    type Format<'a> = FormatRefWithRule<'a, TypeParam, FormatTypeParam, PyFormatContext<'ast>>;
+
+    fn format(&self) -> Self::Format<'_> {
+        FormatRefWithRule::new(self, FormatTypeParam)
+    }
+}
+
+impl<'ast> IntoFormat<PyFormatContext<'ast>> for TypeParam {
+    type Format = FormatOwnedWithRule<TypeParam, FormatTypeParam, PyFormatContext<'ast>>;
+
+    fn into_format(self) -> Self::Format {
+        FormatOwnedWithRule::new(self, FormatTypeParam)
+    }
+}

crates/ruff_python_formatter/src/type_param/type_param_clause.rs

@@ -0,0 +1,29 @@
+use crate::builders::PyFormatterExtensions;
+use crate::context::PyFormatContext;
+use crate::expression::parentheses::parenthesized;
+use crate::prelude::*;
+use ruff_formatter::{Format, FormatResult};
+use ruff_python_ast::TypeParam;
+use ruff_text_size::TextSize;
+
+pub(crate) struct FormatTypeParamsClause<'a> {
+    pub(crate) sequence_end: TextSize,
+    pub(crate) type_params: &'a Vec<TypeParam>,
+}
+
+/// Formats a sequence of [`TypeParam`] nodes.
+impl Format<PyFormatContext<'_>> for FormatTypeParamsClause<'_> {
+    fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
+        if self.type_params.is_empty() {
+            return Ok(());
+        }
+
+        let items = format_with(|f| {
+            f.join_comma_separated(self.sequence_end)
+                .nodes(self.type_params.iter())
+                .finish()
+        });
+
+        parenthesized("[", &items, "]").fmt(f)
+    }
+}

crates/ruff_python_formatter/src/type_param/type_param_param_spec.rs

@@ -0,0 +1,14 @@
+use crate::{AsFormat, FormatNodeRule, PyFormatter};
+use ruff_formatter::prelude::text;
+use ruff_formatter::{write, Buffer, FormatResult};
+use ruff_python_ast::TypeParamParamSpec;
+
+#[derive(Default)]
+pub struct FormatTypeParamParamSpec;
+
+impl FormatNodeRule<TypeParamParamSpec> for FormatTypeParamParamSpec {
+    fn fmt_fields(&self, item: &TypeParamParamSpec, f: &mut PyFormatter) -> FormatResult<()> {
+        let TypeParamParamSpec { range: _, name } = item;
+        write!(f, [text("**"), name.format()])
+    }
+}

crates/ruff_python_formatter/src/type_param/type_param_type_var.rs

@@ -0,0 +1,22 @@
+use crate::{AsFormat, FormatNodeRule, PyFormatter};
+use ruff_formatter::prelude::{space, text};
+use ruff_formatter::{write, Buffer, Format, FormatResult};
+use ruff_python_ast::TypeParamTypeVar;
+
+#[derive(Default)]
+pub struct FormatTypeParamTypeVar;
+
+impl FormatNodeRule<TypeParamTypeVar> for FormatTypeParamTypeVar {
+    fn fmt_fields(&self, item: &TypeParamTypeVar, f: &mut PyFormatter) -> FormatResult<()> {
+        let TypeParamTypeVar {
+            range: _,
+            name,
+            bound,
+        } = item;
+        name.format().fmt(f)?;
+        if let Some(bound) = bound {
+            write!(f, [text(":"), space(), bound.format()])?;
+        }
+        Ok(())
+    }
+}

crates/ruff_python_formatter/src/type_param/type_param_type_var_tuple.rs

@@ -0,0 +1,14 @@
+use crate::{AsFormat, FormatNodeRule, PyFormatter};
+use ruff_formatter::prelude::text;
+use ruff_formatter::{write, Buffer, FormatResult};
+use ruff_python_ast::TypeParamTypeVarTuple;
+
+#[derive(Default)]
+pub struct FormatTypeParamTypeVarTuple;
+
+impl FormatNodeRule<TypeParamTypeVarTuple> for FormatTypeParamTypeVarTuple {
+    fn fmt_fields(&self, item: &TypeParamTypeVarTuple, f: &mut PyFormatter) -> FormatResult<()> {
+        let TypeParamTypeVarTuple { range: _, name } = item;
+        write!(f, [text("*"), name.format()])
+    }
+}

crates/ruff_python_formatter/src/type_param/type_params.rs

@@ -0,0 +1,41 @@
+use crate::builders::PyFormatterExtensions;
+use crate::comments::trailing_comments;
+use crate::expression::parentheses::parenthesized;
+use crate::prelude::*;
+use ruff_formatter::write;
+use ruff_formatter::FormatResult;
+use ruff_python_ast::node::AstNode;
+
+use ruff_python_ast::TypeParams;
+
+#[derive(Default)]
+pub struct FormatTypeParams;
+
+/// Formats a sequence of [`TypeParam`] nodes.
+impl FormatNodeRule<TypeParams> for FormatTypeParams {
+    fn fmt_fields(&self, item: &TypeParams, f: &mut PyFormatter) -> FormatResult<()> {
+        // A dangling comment indicates a comment on the same line as the opening bracket, e.g.:
+        // ```python
+        // type foo[  # This type parameter clause has a dangling comment.
+        //     a,
+        //     b,
+        //     c,
+        // ] = ...
+        let comments = f.context().comments().clone();
+        let dangling_comments = comments.dangling_comments(item.as_any_node_ref());
+        write!(f, [trailing_comments(dangling_comments)])?;
+
+        let items = format_with(|f| {
+            f.join_comma_separated(item.range.end())
+                .nodes(item.type_params.iter())
+                .finish()
+        });
+
+        parenthesized("[", &items, "]").fmt(f)
+    }
+
+    fn fmt_dangling_comments(&self, _node: &TypeParams, _f: &mut PyFormatter) -> FormatResult<()> {
+        // Handled in `fmt_fields`
+        Ok(())
+    }
+}