Format named expressions (walrus operator) (#5642)

## 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

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

@@ -0,0 +1,13 @@
+y = 1
+
+if (
+    # 1
+    x # 2
+    := # 3
+    y # 4
+):
+    pass
+
+y0 = (y1 := f(x))
+
+f(x:=y, z=True)

crates/ruff_python_formatter/src/expression/expr_named_expr.rs

@@ -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;
 
@@ -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(),
+            ]
+        )
     }
 }
 
@@ -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,
+        }
     }
 }

crates/ruff_python_formatter/tests/snapshots/black_compatibility@py_310__pep_572_py310.py.snap

@@ -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[a := 0]
-+x[NOT_YET_IMPLEMENTED_ExprNamedExpr, NOT_YET_IMPLEMENTED_ExprNamedExpr]
++x[a := 0, b := 1]
-+x[5, NOT_YET_IMPLEMENTED_ExprNamedExpr]
++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):
@@ -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[a := 0]
-x[NOT_YET_IMPLEMENTED_ExprNamedExpr, NOT_YET_IMPLEMENTED_ExprNamedExpr]
+x[a := 0, b := 1]
-x[5, NOT_YET_IMPLEMENTED_ExprNamedExpr]
+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 [])):

crates/ruff_python_formatter/tests/snapshots/black_compatibility@py_38__pep_572.py.snap

@@ -59,36 +59,20 @@ while x := f(x):
 ```diff
 --- Black
 +++ Ruff
-@@ -1,47 +1,47 @@
+@@ -2,10 +2,10 @@
--(a := 1)
+ (a := a)
--(a := a)
+ if (match := pattern.search(data)) is None:
--if (match := pattern.search(data)) is None:
-+(NOT_YET_IMPLEMENTED_ExprNamedExpr)
-+(NOT_YET_IMPLEMENTED_ExprNamedExpr)
-+if (NOT_YET_IMPLEMENTED_ExprNamedExpr) 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)
+ (y := f(x))
-+y0 = NOT_YET_IMPLEMENTED_ExprNamedExpr
+ y0 = (y1 := f(x))
-+foo(x=(NOT_YET_IMPLEMENTED_ExprNamedExpr))
+ foo(x=(y := f(x)))
- 
+@@ -19,29 +19,29 @@
- 
--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):
      pass
  
  
@@ -96,99 +80,89 @@ while x := f(x):
 -(x := lambda: 1)
 -(x := lambda: (y := 1))
 -lambda line: (m := re.match(pattern, line)) and m.group(1)
--x = (y := 0)
++lambda x: True
--(z := (y := (x := 0)))
++(x := lambda x: True)
++(x := lambda x: True)
++lambda x: True
+ x = (y := 0)
+ (z := (y := (x := 0)))
 -(info := (name, phone, *rest))
--(x := 1, 2)
++(info := (name, phone, *NOT_YET_IMPLEMENTED_ExprStarred))
--(total := total + tax)
+ (x := 1, 2)
--len(lines := f.readlines())
+ (total := total + tax)
--foo(x := 3, cat="vector")
+ len(lines := f.readlines())
--foo(cat=(category := "vector"))
+ 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)
-+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))
 +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 (ans := self._check_nans(context=context)):
-+if self._is_special and (NOT_YET_IMPLEMENTED_ExprNamedExpr):
      return ans
--foo(b := 2, a=1)
+ foo(b := 2, a=1)
 -foo((b := 2), a=1)
--foo(c=(b := 2), a=1)
++foo(b := 2, a=1)
-+foo(NOT_YET_IMPLEMENTED_ExprNamedExpr, a=1)
+ foo(c=(b := 2), a=1)
-+foo(NOT_YET_IMPLEMENTED_ExprNamedExpr, a=1)
-+foo(c=(NOT_YET_IMPLEMENTED_ExprNamedExpr), 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)
+(a := 1)
-(NOT_YET_IMPLEMENTED_ExprNamedExpr)
+(a := a)
-if (NOT_YET_IMPLEMENTED_ExprNamedExpr) is None:
+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)
+(y := f(x))
-y0 = NOT_YET_IMPLEMENTED_ExprNamedExpr
+y0 = (y1 := f(x))
-foo(x=(NOT_YET_IMPLEMENTED_ExprNamedExpr))
+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)
+(x := lambda x: True)
-(NOT_YET_IMPLEMENTED_ExprNamedExpr)
+(x := lambda x: True)
 lambda x: True
-x = NOT_YET_IMPLEMENTED_ExprNamedExpr
+x = (y := 0)
-(NOT_YET_IMPLEMENTED_ExprNamedExpr)
+(z := (y := (x := 0)))
-(NOT_YET_IMPLEMENTED_ExprNamedExpr)
+(info := (name, phone, *NOT_YET_IMPLEMENTED_ExprStarred))
-(NOT_YET_IMPLEMENTED_ExprNamedExpr, 2)
+(x := 1, 2)
-(NOT_YET_IMPLEMENTED_ExprNamedExpr)
+(total := total + tax)
-len(NOT_YET_IMPLEMENTED_ExprNamedExpr)
+len(lines := f.readlines())
-foo(NOT_YET_IMPLEMENTED_ExprNamedExpr, cat="vector")
+foo(x := 3, cat="vector")
-foo(cat=(NOT_YET_IMPLEMENTED_ExprNamedExpr))
+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(b := 2, a=1)
-foo(NOT_YET_IMPLEMENTED_ExprNamedExpr, a=1)
+foo(b := 2, a=1)
-foo(c=(NOT_YET_IMPLEMENTED_ExprNamedExpr), 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
 ```
 

crates/ruff_python_formatter/tests/snapshots/black_compatibility@py_39__pep_572_py39.py.snap

@@ -22,15 +22,13 @@ 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[((a := 1), (b := 3))]
-+x[((NOT_YET_IMPLEMENTED_ExprNamedExpr), (NOT_YET_IMPLEMENTED_ExprNamedExpr))]
 ```
 
 ## Ruff Output
@@ -38,11 +36,11 @@ x[(a := 1), (b := 3)]
 ```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[(a := 1)]
-x[((NOT_YET_IMPLEMENTED_ExprNamedExpr), (NOT_YET_IMPLEMENTED_ExprNamedExpr))]
+x[((a := 1), (b := 3))]
 ```
 
 ## Black Output

crates/ruff_python_formatter/tests/snapshots/black_compatibility@py_39__python39.py.snap

@@ -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():
      ...
 ```
@@ -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():
     ...
 ```

crates/ruff_python_formatter/tests/snapshots/black_compatibility@py_39__remove_with_brackets.py.snap

@@ -1,216 +0,0 @@
----
-source: crates/ruff_python_formatter/tests/fixtures.rs
-input_file: crates/ruff_python_formatter/resources/test/fixtures/black/py_39/remove_with_brackets.py
----
-## Input
-
-```py
-with (open("bla.txt")):
-    pass
-
-with (open("bla.txt")), (open("bla.txt")):
-    pass
-
-with (open("bla.txt") as f):
-    pass
-
-# Remove brackets within alias expression
-with (open("bla.txt")) as f:
-    pass
-
-# Remove brackets around one-line context managers
-with (open("bla.txt") as f, (open("x"))):
-    pass
-
-with ((open("bla.txt")) as f, open("x")):
-    pass
-
-with (CtxManager1() as example1, CtxManager2() as example2):
-    ...
-
-# Brackets remain when using magic comma
-with (CtxManager1() as example1, CtxManager2() as example2,):
-    ...
-
-# Brackets remain for multi-line context managers
-with (CtxManager1() as example1, CtxManager2() as example2, CtxManager2() as example2, CtxManager2() as example2, CtxManager2() as example2):
-    ...
-
-# Don't touch assignment expressions
-with (y := open("./test.py")) as f:
-    pass
-
-# Deeply nested examples
-# N.B. Multiple brackets are only possible
-# around the context manager itself.
-# Only one brackets is allowed around the
-# alias expression or comma-delimited context managers.
-with (((open("bla.txt")))):
-    pass
-
-with (((open("bla.txt")))), (((open("bla.txt")))):
-    pass
-
-with (((open("bla.txt")))) as f:
-    pass
-
-with ((((open("bla.txt")))) as f):
-    pass
-
-with ((((CtxManager1()))) as example1, (((CtxManager2()))) as example2):
-    ...
-```
-
-## Black Differences
-
-```diff
---- Black
-+++ Ruff
-@@ -39,7 +39,7 @@
-     ...
- 
- # Don't touch assignment expressions
--with (y := open("./test.py")) as f:
-+with NOT_YET_IMPLEMENTED_ExprNamedExpr as f:
-     pass
- 
- # Deeply nested examples
-```
-
-## Ruff Output
-
-```py
-with open("bla.txt"):
-    pass
-
-with open("bla.txt"), open("bla.txt"):
-    pass
-
-with open("bla.txt") as f:
-    pass
-
-# Remove brackets within alias expression
-with open("bla.txt") as f:
-    pass
-
-# Remove brackets around one-line context managers
-with open("bla.txt") as f, open("x"):
-    pass
-
-with open("bla.txt") as f, open("x"):
-    pass
-
-with CtxManager1() as example1, CtxManager2() as example2:
-    ...
-
-# Brackets remain when using magic comma
-with (
-    CtxManager1() as example1,
-    CtxManager2() as example2,
-):
-    ...
-
-# Brackets remain for multi-line context managers
-with (
-    CtxManager1() as example1,
-    CtxManager2() as example2,
-    CtxManager2() as example2,
-    CtxManager2() as example2,
-    CtxManager2() as example2,
-):
-    ...
-
-# Don't touch assignment expressions
-with NOT_YET_IMPLEMENTED_ExprNamedExpr as f:
-    pass
-
-# Deeply nested examples
-# N.B. Multiple brackets are only possible
-# around the context manager itself.
-# Only one brackets is allowed around the
-# alias expression or comma-delimited context managers.
-with open("bla.txt"):
-    pass
-
-with open("bla.txt"), open("bla.txt"):
-    pass
-
-with open("bla.txt") as f:
-    pass
-
-with open("bla.txt") as f:
-    pass
-
-with CtxManager1() as example1, CtxManager2() as example2:
-    ...
-```
-
-## Black Output
-
-```py
-with open("bla.txt"):
-    pass
-
-with open("bla.txt"), open("bla.txt"):
-    pass
-
-with open("bla.txt") as f:
-    pass
-
-# Remove brackets within alias expression
-with open("bla.txt") as f:
-    pass
-
-# Remove brackets around one-line context managers
-with open("bla.txt") as f, open("x"):
-    pass
-
-with open("bla.txt") as f, open("x"):
-    pass
-
-with CtxManager1() as example1, CtxManager2() as example2:
-    ...
-
-# Brackets remain when using magic comma
-with (
-    CtxManager1() as example1,
-    CtxManager2() as example2,
-):
-    ...
-
-# Brackets remain for multi-line context managers
-with (
-    CtxManager1() as example1,
-    CtxManager2() as example2,
-    CtxManager2() as example2,
-    CtxManager2() as example2,
-    CtxManager2() as example2,
-):
-    ...
-
-# Don't touch assignment expressions
-with (y := open("./test.py")) as f:
-    pass
-
-# Deeply nested examples
-# N.B. Multiple brackets are only possible
-# around the context manager itself.
-# Only one brackets is allowed around the
-# alias expression or comma-delimited context managers.
-with open("bla.txt"):
-    pass
-
-with open("bla.txt"), open("bla.txt"):
-    pass
-
-with open("bla.txt") as f:
-    pass
-
-with open("bla.txt") as f:
-    pass
-
-with CtxManager1() as example1, CtxManager2() as example2:
-    ...
-```
-
-

crates/ruff_python_formatter/tests/snapshots/format@expression__named_expr.py.snap

@@ -0,0 +1,38 @@
+---
+source: crates/ruff_python_formatter/tests/fixtures.rs
+input_file: crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/named_expr.py
+---
+## Input
+```py
+y = 1
+
+if (
+    # 1
+    x # 2
+    := # 3
+    y # 4
+):
+    pass
+
+y0 = (y1 := f(x))
+
+f(x:=y, z=True)
+```
+
+## Output
+```py
+y = 1
+
+if (
+    # 1
+    x := y  # 2  # 3  # 4
+):
+    pass
+
+y0 = (y1 := f(x))
+
+f(x := y, z=True)
+```
+
+
+