From ea956ee7792cca28c9361a6565ba13b9a20913d6 Mon Sep 17 00:00:00 2001
From: leowrites <leo@liuwork.ca>
Date: Sun, 22 Sep 2024 21:58:34 -0400
Subject: [PATCH] Use snapshot(save)

---
 python_ta/debug/snapshot_manager.py | 37 +++++++----------------------
 1 file changed, 9 insertions(+), 28 deletions(-)

diff --git a/python_ta/debug/snapshot_manager.py b/python_ta/debug/snapshot_manager.py
index 04462ad42..c89385554 100644
--- a/python_ta/debug/snapshot_manager.py
+++ b/python_ta/debug/snapshot_manager.py
@@ -33,12 +33,15 @@ def __init__(
 
     def _trace_func(self, frame: types.FrameType, event: str, _arg: Any) -> None:
         if event == "line" and frame.f_locals:
-            # TODO: use snapshot(save=True)
-            var_data = snapshot(include=self.include)
-            json_data = snapshot_to_json(var_data)
-            if json_data:
-                self._output_snapshot(json_data)
-                self.snapshot_counts += 1
+            if self.output_filepath:
+                self.memory_viz_args.extend(
+                    [
+                        "--output",
+                        os.path.join(self.output_filepath, f"snapshot-{self.snapshot_counts}.svg"),
+                    ]
+                )
+            snapshot(include=self.include, save=True, memory_viz_args=self.memory_viz_args)
+            self.snapshot_counts += 1
 
     def get_snapshot_count(self):
         return self.snapshot_counts
@@ -51,25 +54,3 @@ def __enter__(self):
 
     def __exit__(self, exc_type, exc_val, exc_tb) -> None:
         sys.settrace(None)
-
-    def _output_snapshot(self, data):
-        command = ["npx", "memory-viz"]
-        command.extend(self.memory_viz_args)
-        if self.output_filepath is not None:
-            command.extend(
-                [
-                    "--output",
-                    os.path.join(self.output_filepath, f"snapshot-{self.snapshot_counts}.svg"),
-                ]
-            )
-        npx_path = shutil.which("npx")
-        subprocess.run(
-            command,
-            input=json.dumps(data),
-            executable=npx_path,
-            stdout=sys.stdout,
-            stderr=sys.stderr,
-            encoding="utf-8",
-            text=True,
-            check=True,
-        )