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..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,6 +10,10 @@ if some: ... """ +correct_walrus_while_condition = """ +while some := call(): + ... +""" correct_comprehension = """ some = [ @@ -39,6 +43,14 @@ if some := call(): ... """ +wrong_walrus_while_condition = """ +while any(some := call()): + ... +""" +wrong_walrus_while_body = """ +while True: + print(some := call()) +""" @pytest.mark.parametrize( @@ -46,6 +58,7 @@ [ 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( diff --git a/wemake_python_styleguide/violations/consistency.py b/wemake_python_styleguide/violations/consistency.py index dc0ac8a6f..149e5baf3 100644 --- a/wemake_python_styleguide/violations/consistency.py +++ b/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 diff --git a/wemake_python_styleguide/visitors/ast/operators.py b/wemake_python_styleguide/visitors/ast/operators.py index 12990ca2f..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,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))