Implement Snapshot Tracer

.github/workflows/test.yml

@@ -41,8 +41,8 @@ jobs:
           pip install -e .[dev,cfg,z3]
       - name: Run tests
         run: |
-          pytest -vv --cov python_ta --cov-config=.coveragerc --cov-report lcov --ignore=tests/test_type_constraints --ignore=tests/test_type_inference tests --ignore=tests/test_debug/test_accumulation_table.py --ignore=tests/test_debug/test_recursion_table.py
-          pytest -vv tests/test_debug/test_accumulation_table.py tests/test_debug/test_recursion_table.py
+          pytest -vv --cov python_ta --cov-config=.coveragerc --cov-report lcov --ignore=tests/test_type_constraints --ignore=tests/test_type_inference tests --ignore=tests/test_debug/test_accumulation_table.py --ignore=tests/test_debug/test_recursion_table.py --ignore=tests/test_debug/test_snapshot_tracer.py
+          pytest -vv tests/test_debug/test_accumulation_table.py tests/test_debug/test_recursion_table.py tests/test_debug/test_snapshot_tracer.py
       - name: Upload test coverage to coveralls.io
         uses: coverallsapp/github-action@v2.3.0
         env:

.pre-commit-config.yaml

@@ -22,7 +22,8 @@ exclude: |
   (?x)^(
       examples|
       tests/fixtures/sample_dir|
-      tests/test_debug/snapshot_testing_snapshots
+      tests/test_debug/snapshot_testing_snapshots|
+      tests/test_debug/snapshot_tracer_testing_snapshots
   )
 
 ci:

CHANGELOG.md

@@ -11,6 +11,7 @@ and adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 - Added `include_frames` filter to `snapshot`
 - Added `exclude_vars` filter to `snapshot`
+- Added new `python_ta.debug` module with an `SnapshotTracer` context manager for generating memory models
 
 ### 💫 New checkers
 

docs/debug/index.md

@@ -1,8 +1,11 @@
-# Loop and Recursion Tracing
+# Loop, Recursion and Memory Tracing
 
-This page describes an additional PythonTA feature: print-based debugging for loops and recursion.
-This feature makes tracing easier by printing the state of each loop iteration or recursive function call in a nicely-formatted table using the [tabulate] library.
-This functionality is found in the `python_ta.debug` submodule.
+This page describes two additional PythonTA features:
+
+1. Print-based debugging for loops and recursion: This feature makes tracing easier by printing the state of each loop iteration or recursive function call in a nicely-formatted table using the [tabulate] library.
+2. Snapshot-based debugging for visualizing the memory diagram: This feature makes it easier to visualize the Python memory model by leveraging the [MemoryViz](https://github.com/david-yz-liu/memory-viz) library.
+
+These additional features are found in the `python.debug` submodule.
 
 ## Loop tracing with `AccumulationTable`
 
@@ -250,3 +253,50 @@ The `RecursionTable` is a new PythonTA feature and currently has the following k
 
 [tabulate]: https://github.com/astanin/python-tabulate
 [`sys.settrace`]: https://docs.python.org/3/library/sys.html#sys.settrace
+
+## Tracing the Python Memory Model
+
+The following section will focus on tracing the Python memory model. This feature uses the `python_ta.debug.SnapshotTracer` as a context manager to visualize program memory.
+
+### Example usage
+
+```python
+# demo.py
+from python_ta.debug import SnapshotTracer
+
+
+def func_multi_line(output_directory: str = None) -> None:
+    """
+    Function for testing SnapshotTracer
+    """
+    with SnapshotTracer(
+        output_directory=output_directory,
+        include_frames=("func_multi_line",),
+        exclude_vars=("output_directory",),
+        memory_viz_args=MEMORY_VIZ_ARGS,
+    ):
+        num = 123
+        some_string = "Hello, world"
+        num2 = 321
+        arr = [some_string, "string 123321"]
+
+
+if __name__ == '__main__':
+    func_multi_line()
+```
+
+When this function runs, the variables within `func_multi_line` are captured, and memory models are outputted to `output_directory` for each line of code. For the expected output, refer to the snapshots in `tests/test_debug/snapshot_tracer_testing_snapshots/func_multi_line`.
+
+### API
+
+```{eval-rst}
+.. automethod:: python_ta.debug.SnapshotTracer.__init__
+```
+
+### Current Limitations
+
+The `SnapshotTracer` has the following limitations:
+
+1. Due to differences in Python interpreters, this context manager only works with Python versions >= 3.10.
+2. The context manager does not step into any function calls. Calling functions within the traced function may lead to undefined behavior.
+3. `SnapshotTracer` uses [`sys.settrace`] to update variable states, and therefore is not compatible with other libraries (e.g., debuggers, code coverage tools).

python_ta/debug/__init__.py

@@ -1,2 +1,3 @@
 from .accumulation_table import AccumulationTable
 from .recursion_table import RecursionTable
+from .snapshot_tracer import SnapshotTracer

python_ta/debug/snapshot_tracer.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+import copy
+import inspect
+import logging
+import os
+import sys
+import types
+from typing import Any, Optional
+
+from .snapshot import snapshot
+
+
+class SnapshotTracer:
+    """
+    A class used for snapshot-based debugging to visualize program memory at each line in the calling function.
+
+    Instance attributes:
+        output_directory: The directory where the memory model diagrams will be saved. Defaults to the current directory.
+        _snapshot_args: A dictionary of keyword arguments to pass to the `snapshot` function.
+        _snapshot_counts: The number of snapshots taken.
+    """
+
+    output_directory: Optional[str]
+    _snapshot_args: dict[str, Any]
+    _snapshot_counts: int
+
+    def __init__(self, output_directory: Optional[str] = None, **kwargs) -> None:
+        """Initialize a context manager for snapshot-based debugging.
+
+        Args:
+            output_directory: The directory to save the snapshots, defaulting to the current directory.
+                **Note**: Use this argument instead of the `--output` flag in `memory_viz_args` to specify the output directory.
+            **kwargs: All other keyword arguments are passed to `python.debug.snapshot`. Refer to the `snapshot` function for more details.
+        """
+        if sys.version_info < (3, 10, 0):
+            logging.warning("You need Python 3.10 or later to use SnapshotTracer.")
+        if any("--output" in arg for arg in kwargs.get("memory_viz_args", [])):
+            raise ValueError(
+                "Use the output_directory parameter to specify a different output path."
+            )
+        self._snapshot_args = kwargs
+        self._snapshot_args["memory_viz_args"] = copy.deepcopy(kwargs.get("memory_viz_args", []))
+        self._snapshot_counts = 0
+        self.output_directory = output_directory if output_directory else "."
+
+    def _trace_func(self, frame: types.FrameType, event: str, _arg: Any) -> None:
+        """Take a snapshot of the variables in the functions specified in `self.include`"""
+        if event == "line" and frame.f_locals:
+            self._snapshot_args["memory_viz_args"].extend(
+                [
+                    "--output",
+                    os.path.join(
+                        os.path.abspath(self.output_directory),
+                        f"snapshot-{self._snapshot_counts}.svg",
+                    ),
+                ]
+            )
+            snapshot(
+                save=True,
+                **self._snapshot_args,
+            )
+            self._snapshot_counts += 1
+
+    def __enter__(self):
+        """Set up the trace function to take snapshots at each line of code."""
+        func_frame = inspect.getouterframes(inspect.currentframe())[1].frame
+        func_frame.f_trace = self._trace_func
+        sys.settrace(lambda *_args: None)
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Remove the trace function."""
+        sys.settrace(None)
+        inspect.getouterframes(inspect.currentframe())[1].frame.f_trace = None