Fix formatting of lambda star arguments (#6257)

## Summary
Previously, the ruff formatter was removing the star argument of
`lambda` expressions when formatting.

Given the following code snippet
```python
lambda *a: ()
lambda **b: ()
```
it would be formatted to
```python
lambda: ()
lambda: ()
```

We fix this by checking for the presence of `args`, `vararg` or `kwarg`
in the `lambda` expression, before we were only checking for the
presence of `args`.

Fixes #5894

## Test Plan

Add new tests cases.

---------

Co-authored-by: Charlie Marsh <charlie.r.marsh@gmail.com>
This commit is contained in:
Victor Hugo Gomes 2023-08-02 16:31:20 -03:00 committed by GitHub
parent c362ea7fd4
commit 7c5791fb77
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 212 additions and 242 deletions

View file

@ -22,7 +22,8 @@ impl FormatNodeRule<ExprLambda> for FormatExprLambda {
write!(f, [text("lambda")])?;
if !parameters.args.is_empty() {
if !parameters.args.is_empty() || parameters.vararg.is_some() || parameters.kwarg.is_some()
{
write!(
f,
[

View file

@ -278,9 +278,9 @@ pub(crate) struct ArgumentSeparator {
pub(crate) following_start: TextSize,
}
/// Finds slash and star in `f(a, /, b, *, c)`
/// Finds slash and star in `f(a, /, b, *, c)` or `lambda a, /, b, *, c: 1`.
///
/// Returns slash and star
/// Returns the location of the slash and star separators, if any.
pub(crate) fn find_argument_separators(
contents: &str,
parameters: &Parameters,
@ -347,14 +347,21 @@ pub(crate) fn find_argument_separators(
} else {
let mut tokens = SimpleTokenizer::new(contents, parameters.range).skip_trivia();
let lparen = tokens
.next()
.expect("The function definition can't end here");
debug_assert!(lparen.kind() == SimpleTokenKind::LParen, "{lparen:?}");
let star = tokens
let lparen_or_star = tokens
.next()
.expect("The function definition can't end here");
// In a function definition, the first token should always be a `(`; in a lambda
// definition, it _can't_ be a `(`.
let star = if lparen_or_star.kind == SimpleTokenKind::LParen {
tokens
.next()
.expect("The function definition can't end here")
} else {
lparen_or_star
};
debug_assert!(star.kind() == SimpleTokenKind::Star, "{star:?}");
Some(ArgumentSeparator {
preceding_end: parameters.range.start(),
separator: star.range,