Keep lambda parameters on one line and parenthesize the body if it expands (#21385)

## Summary

This PR makes two changes to our formatting of `lambda` expressions:
1. We now parenthesize the body expression if it expands
2. We now try to keep the parameters on a single line

The latter of these fixes #8179:

Black formatting and this PR's formatting:

```py
def a():
    return b(
        c,
        d,
        e,
        f=lambda self, *args, **kwargs: aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa(
            *args, **kwargs
        ),
    )
```

Stable Ruff formatting

```py
def a():
    return b(
        c,
        d,
        e,
        f=lambda self,
        *args,
        **kwargs: aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa(*args, **kwargs),
    )
```

We don't parenthesize the body expression here because the call to
`aaaa...` has its own parentheses, but adding a binary operator shows
the new parenthesization:

```diff
@@ -3,7 +3,7 @@
         c,
         d,
         e,
-        f=lambda self, *args, **kwargs: aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa(
-            *args, **kwargs
-        ) + 1,
+        f=lambda self, *args, **kwargs: (
+            aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa(*args, **kwargs) + 1
+        ),
     )
```

This is actually a new divergence from Black, which formats this input
like this:

```py
def a():
    return b(
        c,
        d,
        e,
        f=lambda self, *args, **kwargs: aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa(
            *args, **kwargs
        )
        + 1,
    )
```

But I think this is an improvement, unlike the case from #8179.

One other, smaller benefit is that because we now add parentheses to
lambda bodies, we also remove redundant parentheses:

```diff
 @pytest.mark.parametrize(
     "f",
     [
-        lambda x: (x.expanding(min_periods=5).cov(x, pairwise=True)),
-        lambda x: (x.expanding(min_periods=5).corr(x, pairwise=True)),
+        lambda x: x.expanding(min_periods=5).cov(x, pairwise=True),
+        lambda x: x.expanding(min_periods=5).corr(x, pairwise=True),
     ],
 )
 def test_moment_functions_zero_length_pairwise(f):
```

## Test Plan

New tests taken from #8465 and probably a few more I should grab from
the ecosystem results.

---------

Co-authored-by: Micha Reiser <micha@reiser.io>
This commit is contained in:
Brent Westbrook
2025-12-12 12:02:25 -05:00
committed by GitHub
parent d5546508cf
commit 0ebdebddd8
9 changed files with 2991 additions and 81 deletions

View File

@@ -1,4 +1,4 @@
use ruff_formatter::{Argument, Arguments, write};
use ruff_formatter::{Argument, Arguments, format_args, write};
use ruff_text_size::{Ranged, TextRange, TextSize};
use crate::context::{NodeLevel, WithNodeLevel};
@@ -33,20 +33,27 @@ impl<'ast> Format<PyFormatContext<'ast>> for ParenthesizeIfExpands<'_, 'ast> {
{
let mut f = WithNodeLevel::new(NodeLevel::ParenthesizedExpression, f);
write!(
f,
[group(&format_with(|f| {
if_group_breaks(&token("(")).fmt(f)?;
if self.indent {
soft_block_indent(&Arguments::from(&self.inner)).fmt(f)?;
} else {
Arguments::from(&self.inner).fmt(f)?;
}
if_group_breaks(&token(")")).fmt(f)
}))]
)
if self.indent {
let parens_id = f.group_id("indented_parenthesize_if_expands");
group(&format_args![
if_group_breaks(&token("(")),
indent_if_group_breaks(
&format_args![soft_line_break(), &Arguments::from(&self.inner)],
parens_id
),
soft_line_break(),
if_group_breaks(&token(")"))
])
.with_id(Some(parens_id))
.fmt(&mut f)
} else {
group(&format_args![
if_group_breaks(&token("(")),
Arguments::from(&self.inner),
if_group_breaks(&token(")")),
])
.fmt(&mut f)
}
}
}
}

View File

@@ -1,15 +1,21 @@
use ruff_formatter::write;
use ruff_python_ast::AnyNodeRef;
use ruff_python_ast::ExprLambda;
use ruff_formatter::{FormatRuleWithOptions, RemoveSoftLinesBuffer, format_args, write};
use ruff_python_ast::{AnyNodeRef, Expr, ExprLambda};
use ruff_text_size::Ranged;
use crate::comments::dangling_comments;
use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
use crate::builders::parenthesize_if_expands;
use crate::comments::{SourceComment, dangling_comments, leading_comments, trailing_comments};
use crate::expression::parentheses::{
NeedsParentheses, OptionalParentheses, Parentheses, is_expression_parenthesized,
};
use crate::expression::{CallChainLayout, has_own_parentheses};
use crate::other::parameters::ParametersParentheses;
use crate::prelude::*;
use crate::preview::is_parenthesize_lambda_bodies_enabled;
#[derive(Default)]
pub struct FormatExprLambda;
pub struct FormatExprLambda {
layout: ExprLambdaLayout,
}
impl FormatNodeRule<ExprLambda> for FormatExprLambda {
fn fmt_fields(&self, item: &ExprLambda, f: &mut PyFormatter) -> FormatResult<()> {
@@ -20,13 +26,19 @@ impl FormatNodeRule<ExprLambda> for FormatExprLambda {
body,
} = item;
let body = &**body;
let parameters = parameters.as_deref();
let comments = f.context().comments().clone();
let dangling = comments.dangling(item);
let preview = is_parenthesize_lambda_bodies_enabled(f.context());
write!(f, [token("lambda")])?;
if let Some(parameters) = parameters {
// In this context, a dangling comment can either be a comment between the `lambda` the
// Format any dangling comments before the parameters, but save any dangling comments after
// the parameters/after the header to be formatted with the body below.
let dangling_header_comments = if let Some(parameters) = parameters {
// In this context, a dangling comment can either be a comment between the `lambda` and the
// parameters, or a comment between the parameters and the body.
let (dangling_before_parameters, dangling_after_parameters) = dangling
.split_at(dangling.partition_point(|comment| comment.end() < parameters.start()));
@@ -86,7 +98,7 @@ impl FormatNodeRule<ExprLambda> for FormatExprLambda {
// *x: x
// )
// ```
if comments.has_leading(&**parameters) {
if comments.has_leading(parameters) {
hard_line_break().fmt(f)?;
} else {
write!(f, [space()])?;
@@ -95,32 +107,90 @@ impl FormatNodeRule<ExprLambda> for FormatExprLambda {
write!(f, [dangling_comments(dangling_before_parameters)])?;
}
write!(
f,
[parameters
.format()
.with_options(ParametersParentheses::Never)]
)?;
write!(f, [token(":")])?;
if dangling_after_parameters.is_empty() {
write!(f, [space()])?;
// Try to keep the parameters on a single line, unless there are intervening comments.
if preview && !comments.contains_comments(parameters.into()) {
let mut buffer = RemoveSoftLinesBuffer::new(f);
write!(
buffer,
[parameters
.format()
.with_options(ParametersParentheses::Never)]
)?;
} else {
write!(f, [dangling_comments(dangling_after_parameters)])?;
write!(
f,
[parameters
.format()
.with_options(ParametersParentheses::Never)]
)?;
}
dangling_after_parameters
} else {
write!(f, [token(":")])?;
dangling
};
// In this context, a dangling comment is a comment between the `lambda` and the body.
if dangling.is_empty() {
write!(f, [space()])?;
} else {
write!(f, [dangling_comments(dangling)])?;
}
write!(f, [token(":")])?;
if dangling_header_comments.is_empty() {
write!(f, [space()])?;
} else if !preview {
write!(f, [dangling_comments(dangling_header_comments)])?;
}
write!(f, [body.format()])
if !preview {
return body.format().fmt(f);
}
let fmt_body = FormatBody {
body,
dangling_header_comments,
};
match self.layout {
ExprLambdaLayout::Assignment => fits_expanded(&fmt_body).fmt(f),
ExprLambdaLayout::Default => fmt_body.fmt(f),
}
}
}
#[derive(Debug, Default, Copy, Clone)]
pub enum ExprLambdaLayout {
#[default]
Default,
/// The [`ExprLambda`] is the direct child of an assignment expression, so it needs to use
/// `fits_expanded` to prefer parenthesizing its own body before the assignment tries to
/// parenthesize the whole lambda. For example, we want this formatting:
///
/// ```py
/// long_assignment_target = lambda x, y, z: (
/// x + y + z
/// )
/// ```
///
/// instead of either of these:
///
/// ```py
/// long_assignment_target = (
/// lambda x, y, z: (
/// x + y + z
/// )
/// )
///
/// long_assignment_target = (
/// lambda x, y, z: x + y + z
/// )
/// ```
Assignment,
}
impl FormatRuleWithOptions<ExprLambda, PyFormatContext<'_>> for FormatExprLambda {
type Options = ExprLambdaLayout;
fn with_options(mut self, options: Self::Options) -> Self {
self.layout = options;
self
}
}
@@ -137,3 +207,266 @@ impl NeedsParentheses for ExprLambda {
}
}
}
struct FormatBody<'a> {
body: &'a Expr,
/// Dangling comments attached to the lambda header that should be formatted with the body.
///
/// These can include both own-line and end-of-line comments. For lambdas with parameters, this
/// means comments after the parameters:
///
/// ```py
/// (
/// lambda x, y # 1
/// # 2
/// : # 3
/// # 4
/// x + y
/// )
/// ```
///
/// Or all dangling comments for lambdas without parameters:
///
/// ```py
/// (
/// lambda # 1
/// # 2
/// : # 3
/// # 4
/// 1
/// )
/// ```
///
/// In most cases these should formatted within the parenthesized body, as in:
///
/// ```py
/// (
/// lambda: ( # 1
/// # 2
/// # 3
/// # 4
/// 1
/// )
/// )
/// ```
///
/// or without `# 2`:
///
/// ```py
/// (
/// lambda: ( # 1 # 3
/// # 4
/// 1
/// )
/// )
/// ```
dangling_header_comments: &'a [SourceComment],
}
impl Format<PyFormatContext<'_>> for FormatBody<'_> {
fn fmt(&self, f: &mut PyFormatter) -> FormatResult<()> {
let FormatBody {
dangling_header_comments,
body,
} = self;
let body = *body;
let comments = f.context().comments().clone();
let body_comments = comments.leading_dangling_trailing(body);
if !dangling_header_comments.is_empty() {
// Split the dangling header comments into trailing comments formatted with the lambda
// header (1) and leading comments formatted with the body (2, 3, 4).
//
// ```python
// (
// lambda # 1
// # 2
// : # 3
// # 4
// y
// )
// ```
//
// Note that these are split based on their line position rather than using
// `partition_point` based on a range, for example.
let (trailing_header_comments, leading_body_comments) = dangling_header_comments
.split_at(
dangling_header_comments
.iter()
.position(|comment| comment.line_position().is_own_line())
.unwrap_or(dangling_header_comments.len()),
);
// If the body is parenthesized and has its own leading comments, preserve the
// separation between the dangling lambda comments and the body comments. For
// example, preserve this comment positioning:
//
// ```python
// (
// lambda: # 1
// # 2
// ( # 3
// x
// )
// )
// ```
//
// 1 and 2 are dangling on the lambda and emitted first, followed by a hard line
// break and the parenthesized body with its leading comments.
//
// However, when removing 2, 1 and 3 can instead be formatted on the same line:
//
// ```python
// (
// lambda: ( # 1 # 3
// x
// )
// )
// ```
let comments = f.context().comments();
if is_expression_parenthesized(body.into(), comments.ranges(), f.context().source())
&& comments.has_leading(body)
{
trailing_comments(dangling_header_comments).fmt(f)?;
// Note that `leading_body_comments` have already been formatted as part of
// `dangling_header_comments` above, but their presence still determines the spacing
// here.
if leading_body_comments.is_empty() {
space().fmt(f)?;
} else {
hard_line_break().fmt(f)?;
}
body.format().with_options(Parentheses::Always).fmt(f)
} else {
write!(
f,
[
space(),
token("("),
trailing_comments(trailing_header_comments),
block_indent(&format_args!(
leading_comments(leading_body_comments),
body.format().with_options(Parentheses::Never)
)),
token(")")
]
)
}
}
// If the body has comments, we always want to preserve the parentheses. This also
// ensures that we correctly handle parenthesized comments, and don't need to worry
// about them in the implementation below.
else if body_comments.has_leading() || body_comments.has_trailing_own_line() {
body.format().with_options(Parentheses::Always).fmt(f)
}
// Calls and subscripts require special formatting because they have their own
// parentheses, but they can also have an arbitrary amount of text before the
// opening parenthesis. We want to avoid cases where we keep a long callable on the
// same line as the lambda parameters. For example, `db_evmtx...` in:
//
// ```py
// transaction_count = self._query_txs_for_range(
// get_count_fn=lambda from_ts, to_ts, _chain_id=chain_id: db_evmtx.count_transactions_in_range(
// chain_id=_chain_id,
// from_ts=from_ts,
// to_ts=to_ts,
// ),
// )
// ```
//
// should cause the whole lambda body to be parenthesized instead:
//
// ```py
// transaction_count = self._query_txs_for_range(
// get_count_fn=lambda from_ts, to_ts, _chain_id=chain_id: (
// db_evmtx.count_transactions_in_range(
// chain_id=_chain_id,
// from_ts=from_ts,
// to_ts=to_ts,
// )
// ),
// )
// ```
else if matches!(body, Expr::Call(_) | Expr::Subscript(_)) {
let unparenthesized = body.format().with_options(Parentheses::Never);
if CallChainLayout::from_expression(
body.into(),
comments.ranges(),
f.context().source(),
) == CallChainLayout::Fluent
{
parenthesize_if_expands(&unparenthesized).fmt(f)
} else {
let unparenthesized = unparenthesized.memoized();
if unparenthesized.inspect(f)?.will_break() {
expand_parent().fmt(f)?;
}
best_fitting![
// body all flat
unparenthesized,
// body expanded
group(&unparenthesized).should_expand(true),
// parenthesized
format_args![token("("), block_indent(&unparenthesized), token(")")]
]
.fmt(f)
}
}
// For other cases with their own parentheses, such as lists, sets, dicts, tuples,
// etc., we can just format the body directly. Their own formatting results in the
// lambda being formatted well too. For example:
//
// ```py
// lambda xxxxxxxxxxxxxxxxxxxx, yyyyyyyyyyyyyyyyyyyy, zzzzzzzzzzzzzzzzzzzz: [xxxxxxxxxxxxxxxxxxxx, yyyyyyyyyyyyyyyyyyyy, zzzzzzzzzzzzzzzzzzzz]
// ```
//
// gets formatted as:
//
// ```py
// lambda xxxxxxxxxxxxxxxxxxxx, yyyyyyyyyyyyyyyyyyyy, zzzzzzzzzzzzzzzzzzzz: [
// xxxxxxxxxxxxxxxxxxxx,
// yyyyyyyyyyyyyyyyyyyy,
// zzzzzzzzzzzzzzzzzzzz
// ]
// ```
else if has_own_parentheses(body, f.context()).is_some() {
body.format().fmt(f)
}
// Finally, for expressions without their own parentheses, use
// `parenthesize_if_expands` to add parentheses around the body, only if it expands
// across multiple lines. The `Parentheses::Never` here also removes unnecessary
// parentheses around lambda bodies that fit on one line. For example:
//
// ```py
// lambda xxxxxxxxxxxxxxxxxxxx, yyyyyyyyyyyyyyyyyyyy, zzzzzzzzzzzzzzzzzzzz: xxxxxxxxxxxxxxxxxxxx + yyyyyyyyyyyyyyyyyyyy + zzzzzzzzzzzzzzzzzzzz
// ```
//
// is formatted as:
//
// ```py
// lambda xxxxxxxxxxxxxxxxxxxx, yyyyyyyyyyyyyyyyyyyy, zzzzzzzzzzzzzzzzzzzz: (
// xxxxxxxxxxxxxxxxxxxx + yyyyyyyyyyyyyyyyyyyy + zzzzzzzzzzzzzzzzzzzz
// )
// ```
//
// while
//
// ```py
// lambda xxxxxxxxxxxxxxxxxxxx: (xxxxxxxxxxxxxxxxxxxx + 1)
// ```
//
// is formatted as:
//
// ```py
// lambda xxxxxxxxxxxxxxxxxxxx: xxxxxxxxxxxxxxxxxxxx + 1
// ```
else {
parenthesize_if_expands(&body.format().with_options(Parentheses::Never)).fmt(f)
}
}
}

View File

@@ -52,3 +52,10 @@ pub(crate) const fn is_avoid_parens_for_long_as_captures_enabled(
) -> bool {
context.is_preview()
}
/// Returns `true` if the
/// [`parenthesize_lambda_bodies`](https://github.com/astral-sh/ruff/pull/21385) preview style is
/// enabled.
pub(crate) const fn is_parenthesize_lambda_bodies_enabled(context: &PyFormatContext) -> bool {
context.is_preview()
}

View File

@@ -9,6 +9,7 @@ use crate::comments::{
Comments, LeadingDanglingTrailingComments, SourceComment, trailing_comments,
};
use crate::context::{NodeLevel, WithNodeLevel};
use crate::expression::expr_lambda::ExprLambdaLayout;
use crate::expression::parentheses::{
NeedsParentheses, OptionalParentheses, Parentheses, Parenthesize, is_expression_parenthesized,
optional_parentheses,
@@ -18,6 +19,7 @@ use crate::expression::{
maybe_parenthesize_expression,
};
use crate::other::interpolated_string::InterpolatedStringLayout;
use crate::preview::is_parenthesize_lambda_bodies_enabled;
use crate::statement::trailing_semicolon;
use crate::string::StringLikeExtensions;
use crate::string::implicit::{
@@ -303,12 +305,7 @@ impl Format<PyFormatContext<'_>> for FormatStatementsLastExpression<'_> {
&& format_implicit_flat.is_none()
&& format_interpolated_string.is_none()
{
return maybe_parenthesize_expression(
value,
*statement,
Parenthesize::IfBreaks,
)
.fmt(f);
return maybe_parenthesize_value(value, *statement).fmt(f);
}
let comments = f.context().comments().clone();
@@ -586,11 +583,7 @@ impl Format<PyFormatContext<'_>> for FormatStatementsLastExpression<'_> {
space(),
operator,
space(),
maybe_parenthesize_expression(
value,
*statement,
Parenthesize::IfBreaks
)
maybe_parenthesize_value(value, *statement)
]
);
}
@@ -1369,3 +1362,32 @@ fn is_attribute_with_parenthesized_value(target: &Expr, context: &PyFormatContex
_ => false,
}
}
/// Like [`maybe_parenthesize_expression`] but with special handling for lambdas in preview.
fn maybe_parenthesize_value<'a>(
expression: &'a Expr,
parent: AnyNodeRef<'a>,
) -> MaybeParenthesizeValue<'a> {
MaybeParenthesizeValue { expression, parent }
}
struct MaybeParenthesizeValue<'a> {
expression: &'a Expr,
parent: AnyNodeRef<'a>,
}
impl Format<PyFormatContext<'_>> for MaybeParenthesizeValue<'_> {
fn fmt(&self, f: &mut PyFormatter) -> FormatResult<()> {
let MaybeParenthesizeValue { expression, parent } = self;
if is_parenthesize_lambda_bodies_enabled(f.context())
&& let Expr::Lambda(lambda) = expression
&& !f.context().comments().has_leading(lambda)
{
parenthesize_if_expands(&lambda.format().with_options(ExprLambdaLayout::Assignment))
.fmt(f)
} else {
maybe_parenthesize_expression(expression, *parent, Parenthesize::IfBreaks).fmt(f)
}
}
}