Handle while True loops without break statements (#378)

Author: kreathon

CHANGELOG.md

@@ -1,5 +1,6 @@
 # 2.14 (2024-12-08)
 
+* Handle `while True` loops without `break` statements (kreathon).
 * Improve reachability analysis (kreathon, #270, #302).
 * Add type hints for `get_unused_code` and the fields of the `Item` class (John Doknjas, #361).
 

tests/test_reachability.py

@@ -731,13 +731,77 @@ def test_while_true_else(v):
     check_unreachable(v, 4, 1, "else")
 
 
+def test_while_true_no_fall_through(v):
+    v.scan(
+        """\
+while True:
+    raise Exception() 
+print(":-(")
+"""
+    )
+    check_unreachable(v, 3, 1, "while")
+
+
+def test_while_true_no_fall_through_nested(v):
+    v.scan(
+        """\
+while True:
+    if a > 3:
+        raise Exception() 
+    else:
+        pass
+print(":-(")
+"""
+    )
+    check_unreachable(v, 6, 1, "while")
+
+
+def test_while_true_no_fall_through_nested_loops(v):
+    v.scan(
+        """\
+while True:
+    for _ in range(3):
+        break
+    while False:
+        break
+print(":-(")
+"""
+    )
+    check_multiple_unreachable(v, [(4, 2, "while"), (6, 1, "while")])
+
+
+def test_while_true_fall_through(v):
+    v.scan(
+        """\
+while True:
+    break 
+print(":-)")
+"""
+    )
+    assert v.unreachable_code == []
+
+
+def test_while_true_fall_through_nested(v):
+    v.scan(
+        """\
+while True:
+    if a > 3:
+        raise Exception() 
+    else:
+        break
+print(":-(")
+"""
+    )
+    assert v.unreachable_code == []
+
+
 def test_while_fall_through(v):
     v.scan(
         """\
 def foo(a):
     while a > 0:
         return 1
-    print(":-(")
+    print(":-)")
 """
     )
     assert v.unreachable_code == []

vulture/reachability.py

@@ -8,25 +8,36 @@ def __init__(self, report):
         self._report = report
         self._no_fall_through_nodes = set()
 
+        # Since we visit the children nodes first, we need to maintain a flag
+        # that indicates if a break statement was seen. When visiting the
+        # parent (While, For, or AsyncFor), the value is checked and reset.
+        # Assumes code is valid (break statements only in loops).
+        self._current_loop_has_break_statement = False
+
     def visit(self, node):
         """When called, all children of this node have already been visited."""
         if isinstance(node, (ast.Break, ast.Continue, ast.Return, ast.Raise)):
             self._mark_as_no_fall_through(node)
+            if isinstance(node, ast.Break):
+                self._current_loop_has_break_statement = True
+
         elif isinstance(
             node,
             (
                 ast.Module,
                 ast.FunctionDef,
                 ast.AsyncFunctionDef,
-                ast.For,
-                ast.AsyncFor,
                 ast.With,
                 ast.AsyncWith,
             ),
         ):
             self._can_fall_through_statements_analysis(node.body)
         elif isinstance(node, ast.While):
             self._handle_reachability_while(node)
+            self._current_loop_has_break_statement = False
+        elif isinstance(node, (ast.For, ast.AsyncFor)):
+            self._can_fall_through_statements_analysis(node.body)
+            self._current_loop_has_break_statement = False
         elif isinstance(node, ast.If):
             self._handle_reachability_if(node)
         elif isinstance(node, ast.IfExp):
@@ -162,6 +173,9 @@ def _handle_reachability_while(self, node):
                     message="unreachable 'else' block",
                 )
 
+            if not self._current_loop_has_break_statement:
+                self._mark_as_no_fall_through(node)
+
         self._can_fall_through_statements_analysis(node.body)
 
     def _handle_reachability_try(self, node):