Skip to content

Commit

Permalink
Format named expressions (walrus operator) (#5642)
Browse files Browse the repository at this point in the history
## Summary

Format named expressions (walrus operator) such a `value := f()`. 

Unlike tuples, named expression parentheses are not part of the range
even when mandatory, so mapping optional parentheses to always gives us
decent formatting without implementing all [PEP
572](https://peps.python.org/pep-0572/) rules on when we need
parentheses where other expressions wouldn't. We might want to revisit
this decision later and implement special cases, but for now this gives
us what we need.

## Test Plan

black fixtures, i added some fixtures and checked django and cpython for
stability.

Closes #5613
  • Loading branch information
konstin authored Jul 10, 2023
1 parent 1e894f3 commit bd8f658
Show file tree
Hide file tree
Showing 8 changed files with 152 additions and 318 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
y = 1

if (
# 1
x # 2
:= # 3
y # 4
):
pass

y0 = (y1 := f(x))

f(x:=y, z=True)
26 changes: 23 additions & 3 deletions crates/ruff_python_formatter/src/expression/expr_named_expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,8 @@ use crate::comments::Comments;
use crate::expression::parentheses::{
default_expression_needs_parentheses, NeedsParentheses, Parentheses, Parenthesize,
};
use crate::{not_yet_implemented, FormatNodeRule, PyFormatter};
use crate::{AsFormat, FormatNodeRule, PyFormatter};
use ruff_formatter::prelude::{space, text};
use ruff_formatter::{write, Buffer, FormatResult};
use rustpython_parser::ast::ExprNamedExpr;

Expand All @@ -11,7 +12,21 @@ pub struct FormatExprNamedExpr;

impl FormatNodeRule<ExprNamedExpr> for FormatExprNamedExpr {
fn fmt_fields(&self, item: &ExprNamedExpr, f: &mut PyFormatter) -> FormatResult<()> {
write!(f, [not_yet_implemented(item)])
let ExprNamedExpr {
target,
value,
range: _,
} = item;
write!(
f,
[
target.format(),
space(),
text(":="),
space(),
value.format(),
]
)
}
}

Expand All @@ -22,6 +37,11 @@ impl NeedsParentheses for ExprNamedExpr {
source: &str,
comments: &Comments,
) -> Parentheses {
default_expression_needs_parentheses(self.into(), parenthesize, source, comments)
match default_expression_needs_parentheses(self.into(), parenthesize, source, comments) {
// Unlike tuples, named expression parentheses are not part of the range even when
// mandatory. See [PEP 572](https://peps.python.org/pep-0572/) for details.
Parentheses::Optional => Parentheses::Always,
parentheses => parentheses,
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -32,9 +32,9 @@ f(x, (a := b + c for c in range(10)), y=z, **q)
-x[a:=0]
-x[a:=0, b:=1]
-x[5, b:=0]
+x[NOT_YET_IMPLEMENTED_ExprNamedExpr]
+x[NOT_YET_IMPLEMENTED_ExprNamedExpr, NOT_YET_IMPLEMENTED_ExprNamedExpr]
+x[5, NOT_YET_IMPLEMENTED_ExprNamedExpr]
+x[a := 0]
+x[a := 0, b := 1]
+x[5, b := 0]
# Walruses are allowed inside generator expressions on function calls since 3.10.
-if any(match := pattern_error.match(s) for s in buffer):
Expand Down Expand Up @@ -62,9 +62,9 @@ f(x, (a := b + c for c in range(10)), y=z, **q)

```py
# Unparenthesized walruses are now allowed in indices since Python 3.10.
x[NOT_YET_IMPLEMENTED_ExprNamedExpr]
x[NOT_YET_IMPLEMENTED_ExprNamedExpr, NOT_YET_IMPLEMENTED_ExprNamedExpr]
x[5, NOT_YET_IMPLEMENTED_ExprNamedExpr]
x[a := 0]
x[a := 0, b := 1]
x[5, b := 0]
# Walruses are allowed inside generator expressions on function calls since 3.10.
if any((NOT_YET_IMPLEMENTED_generator_key for NOT_YET_IMPLEMENTED_generator_key in [])):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -59,136 +59,110 @@ while x := f(x):
```diff
--- Black
+++ Ruff
@@ -1,47 +1,47 @@
-(a := 1)
-(a := a)
-if (match := pattern.search(data)) is None:
+(NOT_YET_IMPLEMENTED_ExprNamedExpr)
+(NOT_YET_IMPLEMENTED_ExprNamedExpr)
+if (NOT_YET_IMPLEMENTED_ExprNamedExpr) is None:
@@ -2,10 +2,10 @@
(a := a)
if (match := pattern.search(data)) is None:
pass
-if match := pattern.search(data):
+if NOT_YET_IMPLEMENTED_ExprNamedExpr:
+if (match := pattern.search(data)):
pass
-[y := f(x), y**2, y**3]
[y := f(x), y**2, y**3]
-filtered_data = [y for x in data if (y := f(x)) is None]
-(y := f(x))
-y0 = (y1 := f(x))
-foo(x=(y := f(x)))
+[NOT_YET_IMPLEMENTED_ExprNamedExpr, y**2, y**3]
+filtered_data = [NOT_YET_IMPLEMENTED_generator_key for NOT_YET_IMPLEMENTED_generator_key in []]
+(NOT_YET_IMPLEMENTED_ExprNamedExpr)
+y0 = NOT_YET_IMPLEMENTED_ExprNamedExpr
+foo(x=(NOT_YET_IMPLEMENTED_ExprNamedExpr))
-def foo(answer=(p := 42)):
+def foo(answer=(NOT_YET_IMPLEMENTED_ExprNamedExpr)):
pass
-def foo(answer: (p := 42) = 5):
+def foo(answer: (NOT_YET_IMPLEMENTED_ExprNamedExpr) = 5):
(y := f(x))
y0 = (y1 := f(x))
foo(x=(y := f(x)))
@@ -19,29 +19,29 @@
pass
-lambda: (x := 1)
-(x := lambda: 1)
-(x := lambda: (y := 1))
-lambda line: (m := re.match(pattern, line)) and m.group(1)
-x = (y := 0)
-(z := (y := (x := 0)))
-(info := (name, phone, *rest))
-(x := 1, 2)
-(total := total + tax)
-len(lines := f.readlines())
-foo(x := 3, cat="vector")
-foo(cat=(category := "vector"))
-if any(len(longline := l) >= 100 for l in lines):
+lambda x: True
+(NOT_YET_IMPLEMENTED_ExprNamedExpr)
+(NOT_YET_IMPLEMENTED_ExprNamedExpr)
+(x := lambda x: True)
+(x := lambda x: True)
+lambda x: True
+x = NOT_YET_IMPLEMENTED_ExprNamedExpr
+(NOT_YET_IMPLEMENTED_ExprNamedExpr)
+(NOT_YET_IMPLEMENTED_ExprNamedExpr)
+(NOT_YET_IMPLEMENTED_ExprNamedExpr, 2)
+(NOT_YET_IMPLEMENTED_ExprNamedExpr)
+len(NOT_YET_IMPLEMENTED_ExprNamedExpr)
+foo(NOT_YET_IMPLEMENTED_ExprNamedExpr, cat="vector")
+foo(cat=(NOT_YET_IMPLEMENTED_ExprNamedExpr))
x = (y := 0)
(z := (y := (x := 0)))
-(info := (name, phone, *rest))
+(info := (name, phone, *NOT_YET_IMPLEMENTED_ExprStarred))
(x := 1, 2)
(total := total + tax)
len(lines := f.readlines())
foo(x := 3, cat="vector")
foo(cat=(category := "vector"))
-if any(len(longline := l) >= 100 for l in lines):
+if any((NOT_YET_IMPLEMENTED_generator_key for NOT_YET_IMPLEMENTED_generator_key in [])):
print(longline)
-if env_base := os.environ.get("PYTHONUSERBASE", None):
+if NOT_YET_IMPLEMENTED_ExprNamedExpr:
+if (env_base := os.environ.get("PYTHONUSERBASE", None)):
return env_base
-if self._is_special and (ans := self._check_nans(context=context)):
+if self._is_special and (NOT_YET_IMPLEMENTED_ExprNamedExpr):
if self._is_special and (ans := self._check_nans(context=context)):
return ans
-foo(b := 2, a=1)
foo(b := 2, a=1)
-foo((b := 2), a=1)
-foo(c=(b := 2), a=1)
+foo(NOT_YET_IMPLEMENTED_ExprNamedExpr, a=1)
+foo(NOT_YET_IMPLEMENTED_ExprNamedExpr, a=1)
+foo(c=(NOT_YET_IMPLEMENTED_ExprNamedExpr), a=1)
+foo(b := 2, a=1)
foo(c=(b := 2), a=1)
-while x := f(x):
+while NOT_YET_IMPLEMENTED_ExprNamedExpr:
+while (x := f(x)):
pass
-while x := f(x):
+while NOT_YET_IMPLEMENTED_ExprNamedExpr:
+while (x := f(x)):
pass
```

## Ruff Output

```py
(NOT_YET_IMPLEMENTED_ExprNamedExpr)
(NOT_YET_IMPLEMENTED_ExprNamedExpr)
if (NOT_YET_IMPLEMENTED_ExprNamedExpr) is None:
(a := 1)
(a := a)
if (match := pattern.search(data)) is None:
pass
if NOT_YET_IMPLEMENTED_ExprNamedExpr:
if (match := pattern.search(data)):
pass
[NOT_YET_IMPLEMENTED_ExprNamedExpr, y**2, y**3]
[y := f(x), y**2, y**3]
filtered_data = [NOT_YET_IMPLEMENTED_generator_key for NOT_YET_IMPLEMENTED_generator_key in []]
(NOT_YET_IMPLEMENTED_ExprNamedExpr)
y0 = NOT_YET_IMPLEMENTED_ExprNamedExpr
foo(x=(NOT_YET_IMPLEMENTED_ExprNamedExpr))
(y := f(x))
y0 = (y1 := f(x))
foo(x=(y := f(x)))
def foo(answer=(NOT_YET_IMPLEMENTED_ExprNamedExpr)):
def foo(answer=(p := 42)):
pass
def foo(answer: (NOT_YET_IMPLEMENTED_ExprNamedExpr) = 5):
def foo(answer: (p := 42) = 5):
pass
lambda x: True
(NOT_YET_IMPLEMENTED_ExprNamedExpr)
(NOT_YET_IMPLEMENTED_ExprNamedExpr)
(x := lambda x: True)
(x := lambda x: True)
lambda x: True
x = NOT_YET_IMPLEMENTED_ExprNamedExpr
(NOT_YET_IMPLEMENTED_ExprNamedExpr)
(NOT_YET_IMPLEMENTED_ExprNamedExpr)
(NOT_YET_IMPLEMENTED_ExprNamedExpr, 2)
(NOT_YET_IMPLEMENTED_ExprNamedExpr)
len(NOT_YET_IMPLEMENTED_ExprNamedExpr)
foo(NOT_YET_IMPLEMENTED_ExprNamedExpr, cat="vector")
foo(cat=(NOT_YET_IMPLEMENTED_ExprNamedExpr))
x = (y := 0)
(z := (y := (x := 0)))
(info := (name, phone, *NOT_YET_IMPLEMENTED_ExprStarred))
(x := 1, 2)
(total := total + tax)
len(lines := f.readlines())
foo(x := 3, cat="vector")
foo(cat=(category := "vector"))
if any((NOT_YET_IMPLEMENTED_generator_key for NOT_YET_IMPLEMENTED_generator_key in [])):
print(longline)
if NOT_YET_IMPLEMENTED_ExprNamedExpr:
if (env_base := os.environ.get("PYTHONUSERBASE", None)):
return env_base
if self._is_special and (NOT_YET_IMPLEMENTED_ExprNamedExpr):
if self._is_special and (ans := self._check_nans(context=context)):
return ans
foo(NOT_YET_IMPLEMENTED_ExprNamedExpr, a=1)
foo(NOT_YET_IMPLEMENTED_ExprNamedExpr, a=1)
foo(c=(NOT_YET_IMPLEMENTED_ExprNamedExpr), a=1)
foo(b := 2, a=1)
foo(b := 2, a=1)
foo(c=(b := 2), a=1)
while NOT_YET_IMPLEMENTED_ExprNamedExpr:
while (x := f(x)):
pass
while NOT_YET_IMPLEMENTED_ExprNamedExpr:
while (x := f(x)):
pass
```

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,27 +22,25 @@ x[(a := 1), (b := 3)]
@@ -1,7 +1,7 @@
# Unparenthesized walruses are now allowed in set literals & set comprehensions
# since Python 3.9
-{x := 1, 2, 3}
{x := 1, 2, 3}
-{x4 := x**5 for x in range(7)}
+{NOT_YET_IMPLEMENTED_ExprNamedExpr, 2, 3}
+{NOT_IMPLEMENTED_set_value for value in NOT_IMPLEMENTED_set}
# We better not remove the parentheses here (since it's a 3.10 feature)
-x[(a := 1)]
x[(a := 1)]
-x[(a := 1), (b := 3)]
+x[(NOT_YET_IMPLEMENTED_ExprNamedExpr)]
+x[((NOT_YET_IMPLEMENTED_ExprNamedExpr), (NOT_YET_IMPLEMENTED_ExprNamedExpr))]
+x[((a := 1), (b := 3))]
```

## Ruff Output

```py
# Unparenthesized walruses are now allowed in set literals & set comprehensions
# since Python 3.9
{NOT_YET_IMPLEMENTED_ExprNamedExpr, 2, 3}
{x := 1, 2, 3}
{NOT_IMPLEMENTED_set_value for value in NOT_IMPLEMENTED_set}
# We better not remove the parentheses here (since it's a 3.10 feature)
x[(NOT_YET_IMPLEMENTED_ExprNamedExpr)]
x[((NOT_YET_IMPLEMENTED_ExprNamedExpr), (NOT_YET_IMPLEMENTED_ExprNamedExpr))]
x[(a := 1)]
x[((a := 1), (b := 3))]
```

## Black Output
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -32,14 +32,17 @@ def f():
@relaxed_decorator[0]
def f():
...
@@ -13,8 +12,6 @@
@@ -13,8 +12,10 @@
...
-@extremely_long_variable_name_that_doesnt_fit := complex.expression(
- with_long="arguments_value_that_wont_fit_at_the_end_of_the_line"
-)
+@NOT_YET_IMPLEMENTED_ExprNamedExpr
+@(
+ extremely_long_variable_name_that_doesnt_fit := complex.expression(
+ with_long="arguments_value_that_wont_fit_at_the_end_of_the_line"
+ )
)
def f():
...
```
Expand All @@ -61,7 +64,11 @@ def f():
...
@NOT_YET_IMPLEMENTED_ExprNamedExpr
@(
extremely_long_variable_name_that_doesnt_fit := complex.expression(
with_long="arguments_value_that_wont_fit_at_the_end_of_the_line"
)
)
def f():
...
```
Expand Down
Loading

0 comments on commit bd8f658

Please sign in to comment.