Add z3 option to one iteration checker (#1105)

CHANGELOG.md

@@ -16,6 +16,7 @@ and adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 - Included the name of redundant variable in `E9959 redundant-assignment` message
 - Update to pylint v3.3 and and astroid v3.3. This added support for Python 3.13 and dropped support for Python 3.8.
 - Added a STRICT_NUMERIC_TYPES configuration to `python_ta.contracts` allowing to enable/disable stricter type checking of numeric types
+- Added `z3` option to `one-iteration-checker` to only check for feasible code blocks based on edge z3 constraints
 
 ### 💫 New checkers
 

python_ta/checkers/one_iteration_checker.py

@@ -19,6 +19,17 @@ class OneIterationChecker(BaseChecker):
             '(e.g., by returning or using the "break" keyword).',
         )
     }
+    options = (
+        (
+            "z3",
+            {
+                "default": False,
+                "type": "yn",
+                "metavar": "<y or n>",
+                "help": "Use Z3 to restrict control flow checks to paths that are logically feasible.",
+            },
+        ),
+    )
 
     @only_required_for_messages("one-iteration")
     def visit_for(self, node: nodes.For) -> None:
@@ -53,6 +64,19 @@ def _check_one_iteration(self, node: Union[nodes.For, nodes.While]) -> bool:
         if preds == []:
             return False
 
+        if self.linter.config.z3:
+            # check whether the loop statement is feasible
+            if not start.cfg_block.is_feasible:
+                return False
+            # check whether the loop body is feasible (the loop has at least one iteration)
+            if not any(
+                succ.is_feasible and node.parent_of(succ.target.statements[0])
+                for succ in start.cfg_block.successors
+            ):
+                return False
+            # filter out edges from unfeasible blocks
+            preds = [pred for pred in preds if pred.source.is_feasible]
+
         for pred in preds:
             stmt = pred.source.statements[0]
             if node.parent_of(stmt) and pred.source.reachable:

tests/test_custom_checkers/test_one_iteration_checker.py

@@ -4,6 +4,7 @@
 
 from python_ta.cfg.visitor import CFGVisitor
 from python_ta.checkers.one_iteration_checker import OneIterationChecker
+from python_ta.transforms.z3_visitor import Z3Visitor
 
 
 class TestOneIterationChecker(pylint.testutils.CheckerTestCase):
@@ -439,3 +440,117 @@ def test_while_with_try_except(self):
 
         with self.assertNoMessages():
             self.checker.visit_while(while_node)
+
+
+class TestOneIterationCheckerZ3Option(pylint.testutils.CheckerTestCase):
+    CHECKER_CLASS = OneIterationChecker
+    CONFIG = {"z3": True}
+
+    def test_z3_one_iteration_break_by_precondition(self):
+        src = """
+        def func(x: int) -> int:
+            '''
+            Preconditions:
+                - x > 5
+            '''
+            while x > 0:
+                if x > 3:
+                    break
+                print(x)
+            return x
+        """
+        z3v = Z3Visitor()
+        mod = z3v.visitor.visit(astroid.parse(src))
+        mod.accept(CFGVisitor())
+        while_node = next(mod.nodes_of_class(nodes.While))
+
+        with self.assertAddsMessages(
+            pylint.testutils.MessageTest(
+                msg_id="one-iteration",
+                node=while_node,
+            ),
+            ignore_position=True,
+        ):
+            self.checker.visit_while(while_node)
+
+    def test_z3_multiple_iterations(self):
+        src = """
+        def func(x: int, y: bool) -> int:
+            '''
+            Preconditions:
+                - x > 5
+            '''
+            while x > 0:
+                if x > 3 and y:
+                    break
+                print(x)
+            return x
+        """
+        z3v = Z3Visitor()
+        mod = z3v.visitor.visit(astroid.parse(src))
+        mod.accept(CFGVisitor())
+        while_node = next(mod.nodes_of_class(nodes.While))
+
+        with self.assertNoMessages():
+            self.checker.visit_while(while_node)
+
+    def test_z3_one_iteration_for_loop(self):
+        src = """
+         def func(x: float) -> int:
+            '''
+            Preconditions:
+                - x in [1.0, 2.0, 3.0]
+            '''
+            for i in range(0, 3):
+                if x == 3.0:
+                    break
+            return x
+        """
+        z3v = Z3Visitor()
+        mod = z3v.visitor.visit(astroid.parse(src))
+        mod.accept(CFGVisitor())
+        for_node = next(mod.nodes_of_class(nodes.For))
+
+        with self.assertNoMessages():
+            self.checker.visit_while(for_node)
+
+    def test_z3_one_iteration_unfeasible_loop_body(self):
+        src = """
+        def func(x: int):
+            '''
+            Preconditions:
+                - x > 10
+            '''
+            while x < 0:
+                print("unfeasible")
+                break
+            print("end")
+        """
+        z3v = Z3Visitor()
+        mod = z3v.visitor.visit(astroid.parse(src))
+        mod.accept(CFGVisitor())
+        while_node = next(mod.nodes_of_class(nodes.While))
+
+        with self.assertNoMessages():
+            self.checker.visit_while(while_node)
+
+    def test_z3_one_iteration_unfeasible_loop_statement(self):
+        src = """
+        def func(x: int):
+            '''
+            Preconditions:
+                - x > 10
+            '''
+            if x < 0:
+                while x > 0:
+                    print("unfeasible")
+                    break
+            print("end")
+        """
+        z3v = Z3Visitor()
+        mod = z3v.visitor.visit(astroid.parse(src))
+        mod.accept(CFGVisitor())
+        while_node = next(mod.nodes_of_class(nodes.While))
+
+        with self.assertNoMessages():
+            self.checker.visit_while(while_node)