Add MemoryViz save feature to debug.snapshot function (#1075)

.github/workflows/test.yml

@@ -20,6 +20,10 @@ jobs:
         uses: actions/[email protected]
         with:
           python-version: ${{ matrix.python-version }}
+      - name: Set up Node.js 18
+        uses: actions/setup-node@v4
+        with:
+          node-version: 18
       - name: Set up graphviz
         run: |
           sudo apt-get install -y graphviz

.pre-commit-config.yaml

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

CHANGELOG.md

@@ -21,6 +21,7 @@ and adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 - Stored valid Python function preconditions in initial edge to function code in generated function control flow graphs.
 - Report warning when control flow graph creation encounters a syntax error related to control flow
 - Added autoformat option that runs black formatting tool to python_ta.check_all()
+- Extended the `snapshot` function to optionally generate a svg of the snapshot using MemoryViz when save parameter is true.
 
 ### 💫 New checkers
 

python_ta/debug/snapshot.py

@@ -6,8 +6,15 @@
 from __future__ import annotations
 
 import inspect
+import json
+import logging
+import shutil
+import subprocess
+import sys
 from types import FrameType
-from typing import Any
+from typing import Any, Optional
+
+from packaging.version import Version, parse
 
 
 def get_filtered_global_variables(frame: FrameType) -> dict:
@@ -29,11 +36,21 @@ def get_filtered_global_variables(frame: FrameType) -> dict:
     return {"__main__": true_global_vars}
 
 
-def snapshot():
+def snapshot(
+    save: bool = False,
+    memory_viz_args: Optional[list[str]] = None,
+    memory_viz_version: str = "latest",
+):
     """Capture a snapshot of local variables from the current and outer stack frames
     where the 'snapshot' function is called. Returns a list of dictionaries,
     each mapping function names to their respective local variables.
     Excludes the global module context.
+
+    When save is True, a MemoryViz-created svg is produced.
+    memory_viz_args can be used to pass in options to the MemoryViz CLI.
+    For details on the MemoryViz CLI, see https://www.cs.toronto.edu/~david/memory-viz/docs/cli.
+    memory_viz_version can be used to dictate version, with a default of the latest version.
+    Note that this function is compatible only with MemoryViz version 0.3.1 and above.
     """
     variables = []
     frame = inspect.currentframe().f_back
@@ -47,6 +64,31 @@ def snapshot():
 
         frame = frame.f_back
 
+    if save:
+        json_compatible_vars = snapshot_to_json(variables)
+
+        # Set up command
+        command = ["npx", "memory-viz"]
+        if memory_viz_args:
+            command.extend(memory_viz_args)
+
+        # Ensure valid memory_viz version
+        if memory_viz_version != "latest" and parse(memory_viz_version) < Version("0.3.1"):
+            logging.warning("PythonTA only supports MemoryViz versions 0.3.1 and later.")
+
+        # Create a child to call the MemoryViz CLI
+        npx_path = shutil.which("npx")
+        subprocess.run(
+            command,
+            input=json.dumps(json_compatible_vars),
+            executable=npx_path,
+            stdout=sys.stdout,
+            stderr=sys.stderr,
+            encoding="utf-8",
+            text=True,
+            check=True,
+        )
+
     return variables
 
 

tests/test_debug/snapshot_save_file.py

@@ -0,0 +1,23 @@
+"""
+This Python module is designed for testing the snapshot function's ability to, when save is True,
+create a snapshot svg at the specified file path.
+
+This module is intended exclusively for testing purposes and should not be used for any other purpose.
+"""
+
+import os
+import sys
+
+from python_ta.debug.snapshot import snapshot
+
+test_var1a = "David is cool!"
+test_var2a = "Students Developing Software"
+snapshot(
+    True,
+    [
+        "--output=" + os.path.abspath(sys.argv[1]),
+        "--roughjs-config",
+        "seed=12345",
+    ],
+    "0.3.1",
+)

tests/test_debug/snapshot_save_stdout.py

@@ -0,0 +1,12 @@
+"""
+This Python module is designed for testing the snapshot function's ability to, when save is True,
+return the snapshot svg to stdout.
+
+This module is intended exclusively for testing purposes and should not be used for any other purpose.
+"""
+
+from python_ta.debug.snapshot import snapshot
+
+test_var1a = "David is cool!"
+test_var2a = "Students Developing Software"
+snapshot(True, ["--roughjs-config", "seed=12345"], "0.3.1")

tests/test_debug/test_snapshot.py

@@ -541,3 +541,104 @@ def test_snapshot_to_json_one_class():
     ]
 
     assert json_data == expected_output
+
+
+def test_snapshot_no_save_file():
+    """
+    Tests that snapshot's save feature is not triggered when save = False
+    and ensures no svg files are created.
+    """
+
+    file_path = os.path.join(
+        os.path.dirname(os.path.abspath(__file__)), "snapshot_testing_snapshots"
+    )
+
+    snapshot(
+        False,
+        [
+            "--output=" + file_path,
+            "--roughjs-config",
+            "seed=12345",
+        ],
+    )
+
+    assert not os.path.exists(os.path.join(file_path, "snapshot_testing_snapshots.svg"))
+
+
+def test_snapshot_no_save_stdout(capsys):
+    """
+    Tests that snapshot's save feature is not triggered when save = False
+    """
+    snapshot(False)
+    captured = capsys.readouterr()
+    assert captured.out == ""
+
+
+def test_snapshot_save_create_svg(tmp_path):
+    """
+    Test that snapshot's save feature creates a MemoryViz svg of the stack frame as a file to the specified path.
+    """
+
+    # Calls snapshot in separate file
+    current_directory = os.path.dirname(os.path.abspath(__file__))
+    snapshot_save_path = os.path.join(current_directory, "snapshot_save_file.py")
+    result = subprocess.run(
+        [sys.executable, snapshot_save_path, os.path.abspath(tmp_path)],
+        capture_output=True,
+        text=True,
+        check=True,
+        encoding="utf-8",
+    )
+
+    # Read generated file
+    with open(
+        os.path.join(tmp_path, "test_snapshot_save_create_svg0.svg"),
+        mode="r",
+        encoding="utf-8",
+    ) as gen_svg:
+        generated_svg = gen_svg.read()
+
+    # Read expected file
+    with open(
+        os.path.join(
+            current_directory,
+            "snapshot_testing_snapshots",
+            "snapshot_testing_snapshots_expected.svg",
+        ),
+        mode="r",
+        encoding="utf-8",
+    ) as expected_svg_file:
+        expected_svg = expected_svg_file.read()
+
+    assert generated_svg == expected_svg
+
+
+def test_snapshot_save_stdout():
+    """
+    Test that snapshot's save feature successfully returns a MemoryViz svg of the stack frame to stdout.
+    """
+
+    # Calls snapshot in separate file
+    current_directory = os.path.dirname(os.path.abspath(__file__))
+    snapshot_save_path = os.path.join(current_directory, "snapshot_save_stdout.py")
+    result = subprocess.run(
+        [sys.executable, snapshot_save_path],
+        capture_output=True,
+        encoding="utf-8",
+        text=True,
+        check=True,
+    )
+
+    # Read expected svg file
+    with open(
+        os.path.join(
+            current_directory,
+            "snapshot_testing_snapshots",
+            "snapshot_testing_snapshots_expected_stdout.svg",
+        ),
+        mode="r",
+        encoding="utf-8",
+    ) as expected_svg_file:
+        expected_svg = expected_svg_file.read()
+
+    assert result.stdout == expected_svg