From f4e3cbe9bab2aa2e9bd61098344e8349cfe18d87 Mon Sep 17 00:00:00 2001
From: Raine-Yang-UofT <147955278+Raine-Yang-UofT@users.noreply.github.com>
Date: Tue, 24 Sep 2024 13:00:50 -0400
Subject: [PATCH] Handle annotated and augmented reassignment in Z3 constraint
 parsing (#1088)

---
 CHANGELOG.md                          |  1 +
 python_ta/cfg/graph.py                | 22 +++++++++++------
 tests/test_cfg/test_z3_constraints.py | 34 +++++++++++++++++++++++++++
 3 files changed, 50 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1068a4c05..d23e04291 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -25,6 +25,7 @@ and adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 - Removed `node` attribute for `Z3Parser`
 - Renamed `reduce` method of `Z3Parser` to `parse`
 - Renamed `test_expr_wrapper` to `test_z3_parser`
+- Checked for variable reassignment in `AugAssign` and `AnnAssign` node in parsing edge Z3 constraints
 
 ## [2.8.1] - 2024-08-19
 
diff --git a/python_ta/cfg/graph.py b/python_ta/cfg/graph.py
index df5bfaf5a..da76b3cdb 100644
--- a/python_ta/cfg/graph.py
+++ b/python_ta/cfg/graph.py
@@ -19,9 +19,11 @@
     z3_dependency_available = False
 
 from astroid import (
+    AnnAssign,
     Arguments,
     Assign,
     AssignName,
+    AugAssign,
     Break,
     Continue,
     NodeNG,
@@ -197,7 +199,7 @@ def _get_edges(self, block: CFGBlock, visited: Set[int]) -> Generator[CFGEdge, N
             yield edge
             yield from self._get_edges(edge.target, visited)
 
-    def _get_paths(self) -> List[List[CFGEdge]]:
+    def get_paths(self) -> List[List[CFGEdge]]:
         """Get edges that represent paths from start to end node in depth-first order."""
         paths = []
 
@@ -255,7 +257,7 @@ def update_edge_z3_constraints(self) -> None:
         if not z3_dependency_available:
             return
 
-        for path_id, path in enumerate(self._get_paths()):
+        for path_id, path in enumerate(self.get_paths()):
             # starting a new path
             z3_environment = Z3Environment(self._z3_vars, self.precondition_constraints)
             for edge in path:
@@ -274,11 +276,17 @@ def update_edge_z3_constraints(self) -> None:
 
                 # traverse into target node
                 for node in edge.target.statements:
-                    if isinstance(node, Assign):
-                        # mark reassigned variables
-                        for target in node.targets:
-                            if isinstance(target, AssignName):
-                                z3_environment.assign(target.name)
+                    if isinstance(node, (Assign, AugAssign, AnnAssign)):
+                        self._handle_variable_reassignment(node, z3_environment)
+
+    def _handle_variable_reassignment(self, node: NodeNG, env: Z3Environment) -> None:
+        """Check for reassignment statements and invoke Z3 environment"""
+        if isinstance(node, Assign):
+            for target in node.targets:
+                if isinstance(target, AssignName):
+                    env.assign(target.name)
+        elif isinstance(node, (AugAssign, AnnAssign)):
+            env.assign(node.target.name)
 
 
 class CFGBlock:
diff --git a/tests/test_cfg/test_z3_constraints.py b/tests/test_cfg/test_z3_constraints.py
index 0bfeacd18..c9b6c8850 100644
--- a/tests/test_cfg/test_z3_constraints.py
+++ b/tests/test_cfg/test_z3_constraints.py
@@ -482,6 +482,40 @@ def func(x: float) -> None:
     assert cfg.end.predecessors[0].z3_constraints == {0: []}
 
 
+def test_variable_augumented_reassign() -> None:
+    src = """
+    def func(x: int) -> None:
+        '''
+        Preconditions:
+            - x > 10
+        '''
+        print("initial x:", x)
+        x -= 5
+        print("final x:", x)
+    """
+    cfg = _create_cfg(src, "func")
+    x = z3.Int("x")
+    assert cfg.start.successors[0].z3_constraints == {0: [x > 10]}
+    assert cfg.end.predecessors[0].z3_constraints == {0: []}
+
+
+def test_variable_annotated_reassignment() -> None:
+    src = """
+    def func(x: str) -> None:
+        '''
+        Preconditions:
+            - x[0:2] == "ab"
+        '''
+        print("initial x:", x)
+        x: str = "cd"
+        print("final x:", x)
+    """
+    cfg = _create_cfg(src, "func")
+    x = z3.String("x")
+    assert cfg.start.successors[0].z3_constraints == {0: [z3.SubString(x, 0, 2) == "ab"]}
+    assert cfg.end.predecessors[0].z3_constraints == {0: []}
+
+
 def test_variable_reassignment_in_branch() -> None:
     src = """
     def func(x: float) -> None: