WPS332: allow walrus operator

CHANGELOG.md

@@ -18,6 +18,10 @@ Semantic versioning in our case means:
 
 ## WIP
 
+### Features
+
+- Allows walrus operator in `WPS332`, #3505
+
 ### Bugfixes
 
 - Fixes false positive `WPS457` for ``while True`` loop with ``await`` expressions, #3753

tests/test_visitors/test_ast/test_operators/test_walrus.py

@@ -10,6 +10,10 @@
 if some:
     ...
 """
+correct_walrus_while_condition = """
+while some := call():
+    ...
+"""
 
 correct_comprehension = """
 some = [
@@ -39,13 +43,22 @@
 if some := call():
     ...
 """
+wrong_walrus_while_condition = """
+while any(some := call()):
+    ...
+"""
+wrong_walrus_while_body = """
+while True:
+    print(some := call())
+"""
 
 
 @pytest.mark.parametrize(
     'code',
     [
         correct_assignment,
         correct_if_condition,
+        correct_walrus_while_condition,
         correct_comprehension,
         correct_walrus_comprehension,
         correct_dict_comprehension,
@@ -71,6 +84,8 @@ def test_not_walrus(
     [
         wrong_assignment,
         wrong_if_condition,
+        wrong_walrus_while_condition,
+        wrong_walrus_while_body,
     ],
 )
 def test_walrus(

wemake_python_styleguide/violations/consistency.py

@@ -1306,7 +1306,11 @@ def some_function():
 @final
 class WalrusViolation(ASTViolation):
     """
-    Forbid the use of the walrus operator (`:=`) outside of comprehensions.
+    Forbid the use of the walrus operator (`:=`) in most cases.
+
+    Walrus operator is allowed inside:
+        - comprehensions
+        - top level ``while`` conditions
 
     Reasoning:
         Code with ``:=`` is hardly readable.
@@ -1315,7 +1319,8 @@ class WalrusViolation(ASTViolation):
         Python is not expression-based.
 
     Solution:
-        Avoid using the walrus operator outside comprehensions.
+        Avoid using the walrus operator outside of specific places where it
+        fits.
         Stick to traditional assignment statements for clarity.
 
     Example::
@@ -1335,7 +1340,7 @@ class WalrusViolation(ASTViolation):
 
     """
 
-    error_template = 'Found walrus operator outside a comprehension'
+    error_template = 'Found improper use of a walrus operator'
     code = 332
 
 

wemake_python_styleguide/visitors/ast/operators.py

@@ -4,6 +4,7 @@
 
 from wemake_python_styleguide.compat.aliases import TextNodes
 from wemake_python_styleguide.logic import walk
+from wemake_python_styleguide.logic.nodes import get_parent
 from wemake_python_styleguide.logic.tree.operators import (
     count_unary_operator,
     unwrap_unary_node,
@@ -228,7 +229,7 @@ def _check_string_concat(
 class WalrusVisitor(base.BaseNodeVisitor):
     """We use this visitor to find walrus operators and ban them."""
 
-    _available_parents: ClassVar[AnyNodes] = (
+    _comprehensions: ClassVar[AnyNodes] = (
         ast.ListComp,
         ast.SetComp,
         ast.DictComp,
@@ -239,16 +240,23 @@ def visit_NamedExpr(
         self,
         node: ast.NamedExpr,
     ) -> None:
-        """Disallows walrus ``:=`` operator outside comprehensions."""
-        self._check_walrus_in_comprehesion(node)
+        """Disallows walrus ``:=`` operator in most cases."""
+        self._check_walrus_parent(node)
         self.generic_visit(node)
 
-    def _check_walrus_in_comprehesion(
+    def _check_walrus_parent(
         self,
         node: ast.NamedExpr,
     ) -> None:
-        is_comprension = walk.get_closest_parent(node, self._available_parents)
-        if is_comprension:
+        is_comprehension = walk.get_closest_parent(node, self._comprehensions)
+        if is_comprehension:
+            return
+
+        parent = get_parent(node)
+        is_while_condition = (
+            isinstance(parent, ast.While) and node is parent.test
+        )
+        if is_while_condition:
             return
 
         self.add_violation(consistency.WalrusViolation(node))