Modify SnapshotTracer to generate JSON

packages/python-ta/CHANGELOG.md

@@ -16,6 +16,7 @@ and adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 - Extended `snapshot` to distinguish between nonlocal variables and local variables within a stack frame.
 - Make `watchdog` an optional dependency; users can opt in with `pip install python-ta[watchdog]`. This affects runs of `python_ta.check_all` with the `watch` config option set to `True`.
 - Added `LSPReporter`, a new reporter that outputs lint diagnostics in LSP 3.17-compliant JSON format.
+- Updated `SnapshotTracer` to convert snapshots to JSON with `snapshot_to_json` and pass JSON data to the webstepper template instead of SVG.
 
 ### 💫 New checkers
 

packages/python-ta/src/python_ta/debug/snapshot_tracer.py

@@ -3,7 +3,6 @@
 import base64
 import copy
 import inspect
-import json
 import logging
 import os
 import socket
@@ -15,7 +14,7 @@
 
 from ..util.servers.one_shot_server import open_html_in_browser
 from .id_tracker import IDTracker
-from .snapshot import snapshot
+from .snapshot import snapshot, snapshot_to_json
 
 if TYPE_CHECKING:
     import types
@@ -35,7 +34,7 @@ class SnapshotTracer:
 
     output_directory: Optional[str]
     webstepper: bool
-    _snapshots: list[dict[int, int]]
+    _snapshots: list[dict[str, Any]]
     _snapshot_args: dict[str, Any]
     _first_line: int
 
@@ -77,30 +76,21 @@ def _trace_func(self, frame: types.FrameType, event: str, _arg: Any) -> None:
         if self._first_line == float("inf"):
             self._first_line = frame.f_lineno
         if event == "line":
-            filename = os.path.join(
-                self.output_directory,
-                f"snapshot-{len(self._snapshots)}.svg",
-            )
-            self._snapshot_args["memory_viz_args"].extend(["--output", filename])
-
-            snapshot(
+            snapshot_output = snapshot(
                 id_tracker=self.id_tracker,
-                save=True,
                 **self._snapshot_args,
             )
-
-            self._add_svg_to_map(filename, frame.f_lineno)
-
-    def _add_svg_to_map(self, filename: str, line: int) -> None:
-        """Add the SVG in filename to self._snapshots"""
-        with open(filename) as svg_file:
-            svg_content = svg_file.read()
-            self._snapshots.append(
-                {
-                    "lineNumber": line,
-                    "svg": svg_content,
-                }
-            )
+            json_data = snapshot_to_json(snapshot_output, id_tracker=self.id_tracker)
+            self._add_json_to_map(json_data, frame.f_lineno)
+
+    def _add_json_to_map(self, json_data: list[dict], line: int) -> None:
+        """Add the JSON data to self._snapshots"""
+        self._snapshots.append(
+            {
+                "lineNumber": line,
+                "memoryVizInput": json_data,
+            }
+        )
 
     def __enter__(self):
         """Set up the trace function to take snapshots at each line of code."""
@@ -144,7 +134,7 @@ def _build_self_contained_html(self, func_frame: types.FrameType) -> bytes:
 
         rendered_html = template.render(
             code_text=self._get_code(func_frame),
-            svg_array=self._snapshots,
+            memory_viz_data=self._snapshots,
             bundle_content=bundle_content,
         )
 

packages/python-ta/src/python_ta/debug/webstepper/webstepper_template.html.jinja

@@ -9,7 +9,7 @@
     <div id="root"></div>
     <script>
       window.codeText = `{{ code_text }}`;
-      window.svgArray = {{ svg_array | tojson }};
+      window.memoryVizData = {{ memory_viz_data | tojson }};
     </script>
     <script>
       // Two hacks here:

packages/python-ta/tests/test_debug/test_snapshot_tracer.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import inspect
 import os.path
 import sys
 from typing import Iterator
@@ -27,12 +28,14 @@ def func_one_line(output_directory: str = None) -> None:
     with SnapshotTracer(
         output_directory=output_directory,
         include_frames=(r"^func_one_line$",),
-        exclude_vars=("output_directory",),
+        exclude_vars=("output_directory", "tracer"),
         memory_viz_args=MEMORY_VIZ_ARGS,
         memory_viz_version=MEMORY_VIZ_VERSION,
-    ):
+    ) as tracer:
         num = 123
 
+    return tracer
+
 
 def func_multi_line(output_directory: str = None) -> None:
     """
@@ -41,15 +44,17 @@ def func_multi_line(output_directory: str = None) -> None:
     with SnapshotTracer(
         output_directory=output_directory,
         include_frames=(r"^func_multi_line$",),
-        exclude_vars=("output_directory",),
+        exclude_vars=("output_directory", "tracer"),
         memory_viz_args=MEMORY_VIZ_ARGS,
         memory_viz_version=MEMORY_VIZ_VERSION,
-    ):
+    ) as tracer:
         num = 123
         some_string = "Hello, world"
         num2 = 321
         arr = [some_string, "string 123321"]
 
+    return tracer
+
 
 def func_mutation(output_directory: str = None) -> None:
     """
@@ -58,13 +63,15 @@ def func_mutation(output_directory: str = None) -> None:
     with SnapshotTracer(
         output_directory=output_directory,
         include_frames=(r"^func_mutation$",),
-        exclude_vars=("output_directory",),
+        exclude_vars=("output_directory", "tracer"),
         memory_viz_args=MEMORY_VIZ_ARGS,
         memory_viz_version=MEMORY_VIZ_VERSION,
-    ):
+    ) as tracer:
         num = 123
         num = 321
 
+    return tracer
+
 
 def func_for_loop(output_directory: str = None) -> None:
     """
@@ -73,13 +80,14 @@ def func_for_loop(output_directory: str = None) -> None:
     with SnapshotTracer(
         output_directory=output_directory,
         include_frames=(r"^func_for_loop$",),
-        exclude_vars=("output_directory",),
+        exclude_vars=("output_directory", "tracer"),
         memory_viz_args=MEMORY_VIZ_ARGS,
         memory_viz_version=MEMORY_VIZ_VERSION,
-    ):
+    ) as tracer:
         nums = [1, 2, 3]
         for i in range(len(nums)):
             nums[i] = nums[i] + 1
+    return tracer
 
 
 def func_if_else(output_directory: str = None) -> None:
@@ -89,15 +97,16 @@ def func_if_else(output_directory: str = None) -> None:
     with SnapshotTracer(
         output_directory=output_directory,
         include_frames=(r"^func_if_else$",),
-        exclude_vars=("output_directory",),
+        exclude_vars=("output_directory", "tracer"),
         memory_viz_args=MEMORY_VIZ_ARGS,
         memory_viz_version=MEMORY_VIZ_VERSION,
-    ):
+    ) as tracer:
         num = 10
         if num > 5:
             result = "greater"
         else:
             result = "lesser"
+    return tracer
 
 
 def func_while(output_directory: str = None) -> None:
@@ -107,13 +116,14 @@ def func_while(output_directory: str = None) -> None:
     with SnapshotTracer(
         output_directory=output_directory,
         include_frames=(r"^func_while$",),
-        exclude_vars=("output_directory",),
+        exclude_vars=("output_directory", "tracer"),
         memory_viz_args=MEMORY_VIZ_ARGS,
         memory_viz_version=MEMORY_VIZ_VERSION,
-    ):
+    ) as tracer:
         num = 0
         while num < 3:
             num += 1
+    return tracer
 
 
 def func_no_output_dir() -> None:
@@ -122,10 +132,12 @@ def func_no_output_dir() -> None:
     """
     with SnapshotTracer(
         include_frames=(r"^func_no_output_dir$",),
+        exclude_vars=("tracer",),
         memory_viz_args=MEMORY_VIZ_ARGS,
         memory_viz_version=MEMORY_VIZ_VERSION,
-    ):
+    ) as tracer:
         s = "Hello"
+    return tracer
 
 
 def func_open_webstepper(output_directory: str = None) -> None:
@@ -135,49 +147,34 @@ def func_open_webstepper(output_directory: str = None) -> None:
     with SnapshotTracer(
         output_directory=output_directory,
         include_frames=(r"^func_open_webstepper$",),
-        exclude_vars=("output_directory",),
+        exclude_vars=("output_directory", "tracer"),
         webstepper=True,
         memory_viz_args=MEMORY_VIZ_ARGS,
         memory_viz_version=MEMORY_VIZ_VERSION,
-    ):
+    ) as tracer:
         nums = [1, 2, 3]
         for i in range(len(nums)):
             nums[i] = nums[i] + 1
+    return tracer
 
 
 # Helpers
 
 
-def assert_output_files_match(
-    output_directory: str, snapshot: Snapshot, function_name: str
+def assert_snapshot_data(
+    tracer: SnapshotTracer,
+    expected_num_snapshots: int,
 ) -> None:
     """
-    Assert that the output files in the output directory match the expected output.
+    Assert that SnapshotTracer stored JSON snapshot data correctly.
     """
-    actual_svgs = {}
-    files = os.listdir(output_directory)
-    for file in files:
-        actual_path = os.path.join(output_directory, file)
-        with open(actual_path) as actual_file:
-            actual_svg = actual_file.read()
-            actual_svgs[file] = actual_svg
-    snapshot.assert_match_dir(actual_svgs, function_name)
-
-
-# Fixtures
+    assert len(tracer._snapshots) == expected_num_snapshots
 
+    for snapshot_entry in tracer._snapshots:
+        assert "lineNumber" in snapshot_entry
+        assert "memoryVizInput" in snapshot_entry
 
-@pytest.fixture(scope="function")
-def setup_curr_dir_testing(snapshot: Snapshot) -> Iterator[None]:
-    """
-    Set up and tear down the current directory for the SnapshotTracer tests.
-    """
-    snapshot.snapshot_dir = SNAPSHOT_DIR
-    file_name = "snapshot-0.svg"
-    if os.path.exists(file_name):
-        os.remove(file_name)
-    yield
-    os.remove(file_name)
+        assert isinstance(snapshot_entry["memoryVizInput"], list)
 
 
 # Tests
@@ -204,11 +201,14 @@ def test_snapshot_tracer_with_functions(self, test_func, snapshot, tmp_path):
         """
         Test SnapshotTracer with various simple functions.
         """
-        snapshot.snapshot_dir = SNAPSHOT_DIR
-
-        test_func(str(tmp_path))
+        tracer = test_func(str(tmp_path))
 
-        assert_output_files_match(str(tmp_path), snapshot, test_func.__name__)
+        assert len(tracer._snapshots) > 0
+        for entry in tracer._snapshots:
+            assert "lineNumber" in entry
+            assert "memoryVizInput" in entry
+            assert isinstance(entry["lineNumber"], int)
+            assert isinstance(entry["memoryVizInput"], list)
 
     def test_using_output_flag(self):
         """
@@ -225,18 +225,33 @@ def test_using_output_flag(self):
             ):
                 pass
 
-    def test_no_output_directory(self, snapshot, setup_curr_dir_testing):
+    def test_no_output_directory(self):
         """
         Test SnapshotTracer outputs to the current directory when `output_directory` is not specified.
         """
-        func_no_output_dir()
-
-        with open("snapshot-0.svg") as actual_file:
-            snapshot.assert_match_dir(
-                {"snapshot-0.svg": actual_file.read()}, func_no_output_dir.__name__
-            )
+        tracer = func_no_output_dir()
+        assert len(tracer._snapshots) > 0
 
     def test_serve_html_calls_open_in_browser(self):
         with patch("python_ta.debug.snapshot_tracer.open_html_in_browser") as mock_open:
             func_open_webstepper()
             mock_open.assert_called_once()
+
+    def test_snapshot_contains_json_data(self, tmp_path):
+        tracer = func_multi_line(str(tmp_path))
+        snapshot_entry = tracer._snapshots[0]
+        memory_input = snapshot_entry["memoryVizInput"]
+        assert isinstance(memory_input, list)
+        frame_entries = [entry for entry in memory_input if entry["type"] == ".frame"]
+        assert len(frame_entries) > 0
+
+    def test_snapshot_to_json_called(self, tmp_path):
+        with patch("python_ta.debug.snapshot_tracer.snapshot_to_json") as mock_json:
+            mock_json.return_value = []
+            func_one_line(str(tmp_path))
+            mock_json.assert_called()
+
+    def test_build_html_contains_memoryviz_data(self, tmp_path):
+        tracer = func_one_line(str(tmp_path))
+        assert len(tracer._snapshots) > 0
+        assert all("memoryVizInput" in snap for snap in tracer._snapshots)