From 35184eb12d3a6edf0c03a5d5e9b1f6a8c1c56724 Mon Sep 17 00:00:00 2001 From: fragarie Date: Sun, 1 Feb 2026 17:47:56 +0300 Subject: [PATCH 1/5] WPS332: allow walrus operator --- CHANGELOG.md | 4 ++++ .../test_ast/test_operators/test_walrus.py | 5 +++++ wemake_python_styleguide/violations/consistency.py | 13 ++++++++++--- wemake_python_styleguide/visitors/ast/operators.py | 13 +++++++------ 4 files changed, 26 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2f09b3a66..5e1b2f765 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/tests/test_visitors/test_ast/test_operators/test_walrus.py b/tests/test_visitors/test_ast/test_operators/test_walrus.py index bd6bfd2ea..75540df4e 100644 --- a/tests/test_visitors/test_ast/test_operators/test_walrus.py +++ b/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 = [ @@ -46,6 +50,7 @@ [ correct_assignment, correct_if_condition, + correct_walrus_while_condition, correct_comprehension, correct_walrus_comprehension, correct_dict_comprehension, diff --git a/wemake_python_styleguide/violations/consistency.py b/wemake_python_styleguide/violations/consistency.py index dc0ac8a6f..6c536ea92 100644 --- a/wemake_python_styleguide/violations/consistency.py +++ b/wemake_python_styleguide/violations/consistency.py @@ -1306,7 +1306,7 @@ 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. Reasoning: Code with ``:=`` is hardly readable. @@ -1315,7 +1315,8 @@ class WalrusViolation(ASTViolation): Python is not expression-based. Solution: - Avoid using the walrus operator outside comprehensions. + Avoid using the walrus operator outside comprehensions + or ``while`` conditions. Stick to traditional assignment statements for clarity. Example:: @@ -1325,6 +1326,10 @@ class WalrusViolation(ASTViolation): if some: print(some) + # Correct, but with care + while some := call(): + print(some) + # Wrong: if some := call(): print(some) @@ -1335,7 +1340,9 @@ class WalrusViolation(ASTViolation): """ - error_template = 'Found walrus operator outside a comprehension' + error_template = ( + 'Found walrus operator outside a comprehension or `while` condition' + ) code = 332 diff --git a/wemake_python_styleguide/visitors/ast/operators.py b/wemake_python_styleguide/visitors/ast/operators.py index 12990ca2f..89b5b6158 100644 --- a/wemake_python_styleguide/visitors/ast/operators.py +++ b/wemake_python_styleguide/visitors/ast/operators.py @@ -228,27 +228,28 @@ def _check_string_concat( class WalrusVisitor(base.BaseNodeVisitor): """We use this visitor to find walrus operators and ban them.""" - _available_parents: ClassVar[AnyNodes] = ( + _allowed_parents: ClassVar[AnyNodes] = ( ast.ListComp, ast.SetComp, ast.DictComp, ast.GeneratorExp, + ast.While, ) 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_allowed_parent = walk.get_closest_parent(node, self._allowed_parents) + if is_allowed_parent: return self.add_violation(consistency.WalrusViolation(node)) From f925f7d5f76610e6f24348a3b6b523cb174a2bba Mon Sep 17 00:00:00 2001 From: fragarie Date: Mon, 2 Feb 2026 11:57:08 +0300 Subject: [PATCH 2/5] Clarify docstrings --- .../violations/consistency.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/wemake_python_styleguide/violations/consistency.py b/wemake_python_styleguide/violations/consistency.py index 6c536ea92..b8c658fb2 100644 --- a/wemake_python_styleguide/violations/consistency.py +++ b/wemake_python_styleguide/violations/consistency.py @@ -1308,6 +1308,10 @@ class WalrusViolation(ASTViolation): """ Forbid the use of the walrus operator (`:=`) in most cases. + Walrus operator is allowed inside: + - comprehensions + - ``while`` conditions + Reasoning: Code with ``:=`` is hardly readable. It has big problems with scoping and reading order. @@ -1315,8 +1319,8 @@ class WalrusViolation(ASTViolation): Python is not expression-based. Solution: - Avoid using the walrus operator outside comprehensions - or ``while`` conditions. + Avoid using the walrus operator outside of specific places where it + fits. Stick to traditional assignment statements for clarity. Example:: @@ -1326,10 +1330,6 @@ class WalrusViolation(ASTViolation): if some: print(some) - # Correct, but with care - while some := call(): - print(some) - # Wrong: if some := call(): print(some) @@ -1340,9 +1340,7 @@ class WalrusViolation(ASTViolation): """ - error_template = ( - 'Found walrus operator outside a comprehension or `while` condition' - ) + error_template = 'Found improper use of a walrus operator' code = 332 From 9d38fc9fce504db4b6e250e49e273ead8f0c91bc Mon Sep 17 00:00:00 2001 From: fragarie Date: Mon, 2 Feb 2026 12:45:27 +0300 Subject: [PATCH 3/5] fix walrus ban in while body --- .../test_ast/test_operators/test_walrus.py | 14 ++++++++++++-- wemake_python_styleguide/visitors/ast/operators.py | 14 +++++++++----- 2 files changed, 21 insertions(+), 7 deletions(-) diff --git a/tests/test_visitors/test_ast/test_operators/test_walrus.py b/tests/test_visitors/test_ast/test_operators/test_walrus.py index 75540df4e..103c6d4e6 100644 --- a/tests/test_visitors/test_ast/test_operators/test_walrus.py +++ b/tests/test_visitors/test_ast/test_operators/test_walrus.py @@ -10,10 +10,14 @@ if some: ... """ -correct_walrus_while_condition = """ +correct_walrus_while_condition1 = """ while some := call(): ... """ +correct_walrus_while_condition2 = """ +while any(some := call()): + ... +""" correct_comprehension = """ some = [ @@ -43,6 +47,10 @@ if some := call(): ... """ +wrong_walrus_while_body = """ +while True: + print(some := call()) +""" @pytest.mark.parametrize( @@ -50,7 +58,8 @@ [ correct_assignment, correct_if_condition, - correct_walrus_while_condition, + correct_walrus_while_condition1, + correct_walrus_while_condition2, correct_comprehension, correct_walrus_comprehension, correct_dict_comprehension, @@ -76,6 +85,7 @@ def test_not_walrus( [ wrong_assignment, wrong_if_condition, + wrong_walrus_while_body, ], ) def test_walrus( diff --git a/wemake_python_styleguide/visitors/ast/operators.py b/wemake_python_styleguide/visitors/ast/operators.py index 89b5b6158..31aee1f44 100644 --- a/wemake_python_styleguide/visitors/ast/operators.py +++ b/wemake_python_styleguide/visitors/ast/operators.py @@ -248,8 +248,12 @@ def _check_walrus_parent( self, node: ast.NamedExpr, ) -> None: - is_allowed_parent = walk.get_closest_parent(node, self._allowed_parents) - if is_allowed_parent: - return - - self.add_violation(consistency.WalrusViolation(node)) + allowed_parent = walk.get_closest_parent(node, self._allowed_parents) + if not allowed_parent or ( + isinstance(allowed_parent, ast.While) + and not ( + node is allowed_parent.test + or walk.is_contained_by(node, allowed_parent.test) + ) + ): + self.add_violation(consistency.WalrusViolation(node)) From f771abf3956fef156a0cff77745369a190c1e169 Mon Sep 17 00:00:00 2001 From: fragarie Date: Mon, 2 Feb 2026 14:56:43 +0300 Subject: [PATCH 4/5] forbid walrus in while complex conditions --- .../test_ast/test_operators/test_walrus.py | 14 +++++++------- wemake_python_styleguide/violations/consistency.py | 2 +- wemake_python_styleguide/visitors/ast/operators.py | 5 +---- 3 files changed, 9 insertions(+), 12 deletions(-) diff --git a/tests/test_visitors/test_ast/test_operators/test_walrus.py b/tests/test_visitors/test_ast/test_operators/test_walrus.py index 103c6d4e6..ee854d969 100644 --- a/tests/test_visitors/test_ast/test_operators/test_walrus.py +++ b/tests/test_visitors/test_ast/test_operators/test_walrus.py @@ -10,14 +10,10 @@ if some: ... """ -correct_walrus_while_condition1 = """ +correct_walrus_while_condition = """ while some := call(): ... """ -correct_walrus_while_condition2 = """ -while any(some := call()): - ... -""" correct_comprehension = """ some = [ @@ -47,6 +43,10 @@ if some := call(): ... """ +wrong_walrus_while_condition = """ +while any(some := call()): + ... +""" wrong_walrus_while_body = """ while True: print(some := call()) @@ -58,8 +58,7 @@ [ correct_assignment, correct_if_condition, - correct_walrus_while_condition1, - correct_walrus_while_condition2, + correct_walrus_while_condition, correct_comprehension, correct_walrus_comprehension, correct_dict_comprehension, @@ -85,6 +84,7 @@ def test_not_walrus( [ wrong_assignment, wrong_if_condition, + wrong_walrus_while_condition, wrong_walrus_while_body, ], ) diff --git a/wemake_python_styleguide/violations/consistency.py b/wemake_python_styleguide/violations/consistency.py index b8c658fb2..149e5baf3 100644 --- a/wemake_python_styleguide/violations/consistency.py +++ b/wemake_python_styleguide/violations/consistency.py @@ -1310,7 +1310,7 @@ class WalrusViolation(ASTViolation): Walrus operator is allowed inside: - comprehensions - - ``while`` conditions + - top level ``while`` conditions Reasoning: Code with ``:=`` is hardly readable. diff --git a/wemake_python_styleguide/visitors/ast/operators.py b/wemake_python_styleguide/visitors/ast/operators.py index 31aee1f44..1de288011 100644 --- a/wemake_python_styleguide/visitors/ast/operators.py +++ b/wemake_python_styleguide/visitors/ast/operators.py @@ -251,9 +251,6 @@ def _check_walrus_parent( allowed_parent = walk.get_closest_parent(node, self._allowed_parents) if not allowed_parent or ( isinstance(allowed_parent, ast.While) - and not ( - node is allowed_parent.test - or walk.is_contained_by(node, allowed_parent.test) - ) + and node is not allowed_parent.test ): self.add_violation(consistency.WalrusViolation(node)) From c87ee91990202bc682fad802d7e07501ee602140 Mon Sep 17 00:00:00 2001 From: fragarie Date: Mon, 2 Feb 2026 19:13:34 +0300 Subject: [PATCH 5/5] Split checks for greater readability --- .../visitors/ast/operators.py | 22 ++++++++++++------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/wemake_python_styleguide/visitors/ast/operators.py b/wemake_python_styleguide/visitors/ast/operators.py index 1de288011..a4d7729ea 100644 --- a/wemake_python_styleguide/visitors/ast/operators.py +++ b/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,12 +229,11 @@ def _check_string_concat( class WalrusVisitor(base.BaseNodeVisitor): """We use this visitor to find walrus operators and ban them.""" - _allowed_parents: ClassVar[AnyNodes] = ( + _comprehensions: ClassVar[AnyNodes] = ( ast.ListComp, ast.SetComp, ast.DictComp, ast.GeneratorExp, - ast.While, ) def visit_NamedExpr( @@ -248,9 +248,15 @@ def _check_walrus_parent( self, node: ast.NamedExpr, ) -> None: - allowed_parent = walk.get_closest_parent(node, self._allowed_parents) - if not allowed_parent or ( - isinstance(allowed_parent, ast.While) - and node is not allowed_parent.test - ): - self.add_violation(consistency.WalrusViolation(node)) + 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))