Repository: wemake-services/wemake-python-styleguide

Issue:
InfiniteWhileLoopViolation does not catch anything other than `True`
Currently `InfiniteWhileLoopViolation` does not trigget if you use anything other than `True`.
For example: `1`, `'inf'`, and `[1, 2, 3]` do not count.

We should fix that.

Refs #1857 


Discussion:

Resolving patch:
diff --git a/wemake_python_styleguide/visitors/ast/keywords.py b/wemake_python_styleguide/visitors/ast/keywords.py
--- a/wemake_python_styleguide/visitors/ast/keywords.py
+++ b/wemake_python_styleguide/visitors/ast/keywords.py
@@ -8,7 +8,7 @@
     FunctionNodes,
     TextNodes,
 )
-from wemake_python_styleguide.logic import walk
+from wemake_python_styleguide.logic import walk, walrus
 from wemake_python_styleguide.logic.naming import name_nodes
 from wemake_python_styleguide.logic.nodes import get_parent
 from wemake_python_styleguide.logic.tree import keywords, operators
@@ -405,32 +405,20 @@ class ConstantKeywordVisitor(BaseNodeVisitor):
     )
 
     def visit_While(self, node: ast.While) -> None:
-        """
-        Visits ``while`` keyword and tests that loop will execute.
-
-        Raises:
-            WrongKeywordConditionViolation
-
-        """
+        """Visits ``while`` keyword and tests that loop will execute."""
         self._check_condition(node, node.test)
         self.generic_visit(node)
 
     def visit_Assert(self, node: ast.Assert) -> None:
-        """
-        Visits ``assert`` keyword and tests that condition is correct.
-
-        Raises:
-            WrongKeywordConditionViolation
-
-        """
+        """Visits ``assert`` keyword and tests that condition is correct."""
         self._check_condition(node, node.test)
         self.generic_visit(node)
 
-    def _check_condition(self, node: ast.AST, condition: ast.AST) -> None:
-        real_node = operators.unwrap_unary_node(condition)
-        if isinstance(real_node, ast.NameConstant) and real_node.value is True:
+    def _check_condition(self, node: ast.AST, cond: ast.AST) -> None:
+        if isinstance(cond, ast.NameConstant) and cond.value is True:
             if isinstance(node, ast.While):
-                return  # We should allow `while True:`
+                return  # We should allow plain `while True:`
 
+        real_node = operators.unwrap_unary_node(walrus.get_assigned_expr(cond))
         if isinstance(real_node, self._forbidden_nodes):
-            self.add_violation(WrongKeywordConditionViolation(condition))
+            self.add_violation(WrongKeywordConditionViolation(cond))
diff --git a/wemake_python_styleguide/visitors/ast/loops.py b/wemake_python_styleguide/visitors/ast/loops.py
--- a/wemake_python_styleguide/visitors/ast/loops.py
+++ b/wemake_python_styleguide/visitors/ast/loops.py
@@ -1,5 +1,6 @@
 import ast
 from collections import defaultdict
+from contextlib import suppress
 from typing import ClassVar, DefaultDict, List, Mapping, Sequence, Type, Union
 
 from typing_extensions import final
@@ -123,10 +124,12 @@ def _post_visit(self) -> None:
 class WrongLoopVisitor(base.BaseNodeVisitor):
     """Responsible for examining loops."""
 
-    _breaks: ClassVar[AnyNodes] = (
+    _can_break_loop: ClassVar[AnyNodes] = (
         ast.Break,
         ast.Return,
         ast.Raise,
+        # We only check for `try/except`, not `try/finally`:
+        ast.ExceptHandler,
     )
 
     _containers: ClassVar[_ContainerSpec] = {
@@ -220,9 +223,12 @@ def _check_infinite_while_loop(self, node: AnyLoop) -> None:
         if not isinstance(node, ast.While):
             return
 
-        real_node = operators.unwrap_unary_node(node.test)
-        if isinstance(real_node, ast.NameConstant) and real_node.value is True:
-            if not loops.has_break(node, break_nodes=self._breaks):
+        if loops.has_break(node, break_nodes=self._can_break_loop):
+            return
+
+        with suppress(ValueError):
+            evaled = ast.literal_eval(node.test)
+            if not isinstance(evaled, ast.Name) and bool(evaled):
                 self.add_violation(InfiniteWhileLoopViolation(node))
 
 


Test patch:
diff --git a/tests/test_visitors/test_ast/test_keywords/test_keyword_condition.py b/tests/test_visitors/test_ast/test_keywords/test_keyword_condition.py
--- a/tests/test_visitors/test_ast/test_keywords/test_keyword_condition.py
+++ b/tests/test_visitors/test_ast/test_keywords/test_keyword_condition.py
@@ -1,5 +1,6 @@
 import pytest
 
+from wemake_python_styleguide.compat.constants import PY38
 from wemake_python_styleguide.violations.best_practices import (
     WrongKeywordConditionViolation,
 )
@@ -32,6 +33,7 @@
     '[]',
     '()',
     '{}',  # noqa: P103
+    '+True',
     'False',
     'None',
     '1',
@@ -45,6 +47,14 @@
     '{1 for a in some()}',
     '{"a": a for a in some()}',
     'some if x else other',
+    pytest.param(
+        '(unique := +0)',
+        marks=pytest.mark.skipif(not PY38, reason='walrus appeared in 3.8'),
+    ),
+    pytest.param(
+        '(unique := +True)',
+        marks=pytest.mark.skipif(not PY38, reason='walrus appeared in 3.8'),
+    ),
 ])
 def test_false_condition_keywords(
     assert_errors,
@@ -85,12 +95,7 @@ def test_false_assert_condition_keywords(
     assert_errors(visitor, [WrongKeywordConditionViolation])
 
 
-@pytest.mark.parametrize('code', [
-    while_template,
-    assert_template,
-    assert_with_message_template,
-])
-@pytest.mark.parametrize('condition', [
+fixtures = (
     'call()',
     'name',
     '-name',
@@ -101,8 +106,41 @@ def test_false_assert_condition_keywords(
     'value + 1',
     'x > 1',
     'x + 1 > 1 and y < 1',
+)
+
+
+@pytest.mark.parametrize('code', [
+    while_template,
+])
+@pytest.mark.parametrize('condition', [
+    *fixtures,
+    pytest.param(
+        'x := other()',
+        marks=pytest.mark.skipif(not PY38, reason='walrus appeared in 3.8'),
+    ),
+])
+def test_true_condition_keywords_while(
+    assert_errors,
+    parse_ast_tree,
+    code,
+    condition,
+    default_options,
+):
+    """Testing that true coniditions in keywords are allowed."""
+    tree = parse_ast_tree(code.format(condition))
+
+    visitor = ConstantKeywordVisitor(default_options, tree=tree)
+    visitor.run()
+
+    assert_errors(visitor, [])
+
+
+@pytest.mark.parametrize('code', [
+    assert_template,
+    assert_with_message_template,
 ])
-def test_true_condition_keywords(
+@pytest.mark.parametrize('condition', fixtures)
+def test_true_condition_keywords_assert(
     assert_errors,
     parse_ast_tree,
     code,
diff --git a/tests/test_visitors/test_ast/test_loops/test_loops/test_infinite_while_loops.py b/tests/test_visitors/test_ast/test_loops/test_loops/test_infinite_while_loops.py
--- a/tests/test_visitors/test_ast/test_loops/test_loops/test_infinite_while_loops.py
+++ b/tests/test_visitors/test_ast/test_loops/test_loops/test_infinite_while_loops.py
@@ -19,7 +19,7 @@
 
 template_nested_while2 = """
 while other:
-    while True:
+    while +1:
         {0}
     {0}
 """
@@ -45,11 +45,81 @@ def wrapper():
 
 template_double_while = """
 while True:
-    while True:
+    while 'inf':
         {0}
     {1}
 """
 
+# Correct
+
+correct_while1 = """
+while 1:
+    try:
+       ...
+    except:
+        ...
+"""
+
+correct_while2 = """
+while other:
+    while [1, 2, 3]:
+        try:
+            ...
+        except:
+            ...
+"""
+
+
+correct_while3 = """
+def wrapper():
+    while True:
+        try:
+            ...
+        except:
+            ...
+"""
+
+correct_while4 = """
+while other:
+    ...
+"""
+
+correct_while5 = """
+while 1 + 1:
+    try:
+        ...
+    except:
+        ...
+    finally:
+        ...
+"""
+
+correct_while6 = """
+while 0:
+    ...
+"""
+
+# Do raise:
+
+wrong_while1 = """
+while other:
+    try:
+        ...
+    except:
+        ...
+
+    while True:
+        ...
+"""
+
+wrong_while2 = """
+while True:
+    try:
+        ...
+    finally:
+        ...
+"""
+
 
 @pytest.mark.parametrize('template', [
     template_simple,
@@ -65,7 +135,7 @@ def wrapper():
     'raise Some()',
     'raise',
 ])
-def test_correct_while_loops(
+def test_correct_while_loops_with_statements(
     assert_errors,
     parse_ast_tree,
     keyword,
@@ -81,23 +151,22 @@ def test_correct_while_loops(
     assert_errors(visitor, [])
 
 
-@pytest.mark.parametrize('template', [
-    template_function,
+@pytest.mark.parametrize('code', [
+    correct_while1,
+    correct_while2,
+    correct_while3,
+    correct_while4,
+    correct_while5,
+    correct_while6,
 ])
-@pytest.mark.parametrize('keyword', [
-    'return',
-    'return some',
-])
-def test_correct_while_loops_function(
+def test_correct_while_loops_with_try(
     assert_errors,
     parse_ast_tree,
-    keyword,
-    template,
+    code,
     default_options,
-    mode,
 ):
-    """Testing while loops with ``return`` statements."""
-    tree = parse_ast_tree(mode(template.format(keyword)))
+    """Testing while loops with correct code."""
+    tree = parse_ast_tree(code)
 
     visitor = WrongLoopVisitor(default_options, tree=tree)
     visitor.run()
@@ -106,22 +175,22 @@ def test_correct_while_loops_function(
 
 
 @pytest.mark.parametrize('template', [
-    template_other,
+    template_function,
 ])
 @pytest.mark.parametrize('keyword', [
-    'print(some)',
-    'attr.method()',
-    'a = 1',
+    'return',
+    'return some',
 ])
-def test_other_while_loops(
+def test_correct_while_loops_function(
     assert_errors,
     parse_ast_tree,
     keyword,
     template,
     default_options,
+    mode,
 ):
-    """Testing other while loops with regular code."""
-    tree = parse_ast_tree(template.format(keyword))
+    """Testing while loops with ``return`` statements."""
+    tree = parse_ast_tree(mode(template.format(keyword)))
 
     visitor = WrongLoopVisitor(default_options, tree=tree)
     visitor.run()
@@ -157,6 +226,25 @@ def test_wrong_while_loops(
     assert_errors(visitor, [InfiniteWhileLoopViolation])
 
 
+@pytest.mark.parametrize('code', [
+    wrong_while1,
+    wrong_while2,
+])
+def test_wrong_while_loops_with_try(
+    assert_errors,
+    parse_ast_tree,
+    code,
+    default_options,
+):
+    """Testing while loops with correct code."""
+    tree = parse_ast_tree(code)
+
+    visitor = WrongLoopVisitor(default_options, tree=tree)
+    visitor.run()
+
+    assert_errors(visitor, [InfiniteWhileLoopViolation])
+
+
 @pytest.mark.parametrize('template', [
     template_double_while,
 ])