Augment function CFGs with Z3 vars for parameters (#1056)

.github/workflows/test.yml

@@ -79,6 +79,6 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip setuptools wheel
-          pip install -e .[dev,cfg]
+          pip install -e .[dev,cfg,z3]
       - name: Generate docs
         run: sphinx-build -b html -W docs docs/_build

CHANGELOG.md

@@ -70,6 +70,11 @@ Custom checkers:
 - Fixed minor typo in an error message in `python_ta/cfg/visitor.py`
 - Updated `ExprWrapper` to support `set/list/tuple` literals and `in/not in` operators
 - Updated `snapshot.py` and `test_snapshot.py` to align with MemoryViz 0.2.0 updates
+- Added protected `_z3_vars` attribute to `ControlFlowGraph` to store variables to be used in Z3 solver
+- Removed unused imports from `python_ta/cfg/graph.py`
+- Extended functionality of `ExprWrapper` class to include function definitions' arguments and name assignments
+- Added `z3` to dependencies installed as part of the `docs` job in the GitHub Actions workflow
+- Added tests to maintain/increase coverage of `visitor.py`, `graph.py`, and `ExprWrapper.py`
 
 ## [2.7.0] - 2024-12-14
 

python_ta/cfg/graph.py

@@ -1,8 +1,16 @@
 from __future__ import annotations
 
-from typing import Any, Generator, List, Optional, Set, Tuple, Union
+from typing import Any, Dict, Generator, List, Optional, Set
 
-from astroid import Break, Continue, NodeNG, Raise, Return
+try:
+    from z3 import ExprRef
+
+    from ..transforms import ExprWrapper
+except ImportError:
+    ExprRef = Any
+    ExprWrapper = Any
+
+from astroid import Arguments, Break, Continue, NodeNG, Raise, Return
 
 
 class ControlFlowGraph:
@@ -16,13 +24,27 @@ class ControlFlowGraph:
     block_count: int
     # blocks (with at least one statement) that will never be executed in runtime.
     unreachable_blocks: Set[CFGBlock]
+    # map from variable names to z3 variables
+    _z3_vars: Dict[str, ExprRef]
 
     def __init__(self, cfg_id: int = 0) -> None:
         self.block_count = 0
         self.cfg_id = cfg_id
         self.unreachable_blocks = set()
         self.start = self.create_block()
         self.end = self.create_block()
+        self._z3_vars = {}
+
+    def add_arguments(self, args: Arguments) -> None:
+        self.start.add_statement(args)
+        args.parent.cfg = self
+        args.parent.cfg_block = self.start
+
+        if ExprRef is not Any:
+            # Parse types
+            expr = ExprWrapper(args)
+            z3_vars = expr.parse_arguments(args)
+            self._z3_vars.update(z3_vars)
 
     def create_block(
         self,

python_ta/cfg/visitor.py

@@ -126,9 +126,8 @@ def visit_functiondef(self, func: nodes.FunctionDef) -> None:
             )
         )
 
-        self._current_cfg.start.add_statement(func.args)
-        func.cfg_block = self._current_cfg.start
-        func.cfg = self._current_cfg
+        # Current CFG block is self._current_cfg.start while initially creating the function cfg
+        self._current_cfg.add_arguments(func.args)
 
         preconditions_node = _get_preconditions_node(func)
 

python_ta/transforms/ExprWrapper.py

@@ -3,6 +3,7 @@
 import astroid
 import z3
 from astroid import nodes
+from pylint.checkers.utils import safe_infer
 
 
 class Z3ParseException(Exception):
@@ -26,12 +27,22 @@ class ExprWrapper:
     node: astroid.NodeNG
     types: Dict[str, str]
 
-    def __init__(self, expr: nodes.Expr, types=None):
-        self.node = expr.value
+    def __init__(self, node: astroid.NodeNG, types=None):
         if types is None:
             types = {}
         self.types = types
 
+        if isinstance(node, astroid.Expr):
+            self.node = node.value  # take node attribute to be the value of the expression
+        elif isinstance(node, astroid.Assign):
+            self.node = node.value  # take node attribute as expression (right side) of assignment
+        elif isinstance(node, astroid.Arguments):
+            self.node = node  # take node attribute to be the function declaration node itself
+        else:
+            raise ValueError(
+                "'node' param must be an astroid expression, assignment, or arguments node."
+            )
+
     def reduce(self, node: astroid.NodeNG = None) -> z3.ExprRef:
         """
         Convert astroid node to z3 expression and return it.
@@ -52,6 +63,8 @@ def reduce(self, node: astroid.NodeNG = None) -> z3.ExprRef:
             node = node.value
         elif isinstance(node, nodes.Name):
             node = self.apply_name(node.name)
+        elif isinstance(node, nodes.AssignName):
+            node = self.apply_name(node.name)
         elif isinstance(node, (nodes.List, nodes.Tuple, nodes.Set)):
             node = self.parse_container_op(node)
         else:
@@ -153,7 +166,6 @@ def parse_unary_op(self, node: astroid.UnaryOp) -> z3.ExprRef:
         """Convert an astroid UnaryOp node to a z3 expression."""
         left, op = node.operand, node.op
         left = self.reduce(left)
-
         return self.apply_unary_op(left, op)
 
     def parse_bin_op(self, node: astroid.BinOp) -> z3.ExprRef:
@@ -176,3 +188,24 @@ def parse_container_op(
     ) -> List[z3.ExprRef]:
         """Convert an astroid List, Set, Tuple node to a list of z3 expressions."""
         return [self.reduce(element) for element in node.elts]
+
+    def parse_arguments(self, node: astroid.Arguments) -> Dict[str, z3.ExprRef]:
+        """Convert an astroid Arguments node's parameters to z3 variables."""
+        z3_vars = {}
+
+        annotations = node.annotations
+        arguments = node.args
+        for ann, arg in zip(annotations, arguments):
+            if ann is None:
+                continue
+
+            inferred = safe_infer(ann)
+            if inferred is None or not isinstance(inferred, astroid.ClassDef):
+                continue
+
+            self.types[arg.name] = inferred.name
+
+            if arg.name in self.types and self.types[arg.name] in {"int", "float", "bool"}:
+                z3_vars[arg.name] = self.reduce(arg)
+
+        return z3_vars

python_ta/transforms/__init__.py

@@ -0,0 +1 @@
+from .ExprWrapper import ExprWrapper, Z3ParseException

tests/test_transforms/test_expr_wrapper.py

@@ -0,0 +1,21 @@
+"""Tests for ExprWrapper in python_ta.transforms."""
+
+import pytest
+from astroid import extract_node
+
+from python_ta.transforms import ExprWrapper
+
+
+def test_expr_wrapper_invalid_node() -> None:
+    with pytest.raises(ValueError):
+        ExprWrapper("not a node")
+
+
+def test_expr_wrapper_assignment() -> None:
+    assignment = "n = 2 + 3"
+
+    node = extract_node(assignment)
+    expr = ExprWrapper(node)
+
+    assert expr.node is node.value
+    assert expr.reduce() == 5

tests/test_z3_visitor.py

@@ -4,8 +4,18 @@
 import pytest
 import z3
 
+from python_ta.cfg import ControlFlowGraph
 from python_ta.transforms.z3_visitor import Z3Visitor
 
+# test cases for z3 variables
+z3_vars_example = """
+def f(x: int, y: float, z: bool, a: str):
+    '''
+    This is a function to test z3 vars
+    '''
+    n = x + y - z
+"""
+
 # test cases for arithmetic expressions
 arithmetic_list = [
     """
@@ -168,3 +178,19 @@ def test_constraint(code, expected):
             solver = z3.Solver()
             solver.add(a == e)
             assert solver.check() == z3.sat
+
+
+def test_cfg_z3_vars_initialization():
+    """
+    Test that the cfg's z3 variable mapping is correctly initialized.
+    """
+    node = astroid.extract_node(z3_vars_example)
+
+    cfg = ControlFlowGraph()
+    cfg.add_arguments(node.args)
+
+    # Note that this first assert implicitly includes the assertion that 'a' not in cfg._z3_vars
+    assert len(cfg._z3_vars) == 3
+    assert cfg._z3_vars["x"] == z3.Int("x")
+    assert cfg._z3_vars["y"] == z3.Real("y")
+    assert cfg._z3_vars["z"] == z3.Bool("z")