Add a visualize flag to egraph.saturate for further configuration

egraphs-good · Jan 29, 2025 · e14e29c · e14e29c
1 parent 3829d69
commit e14e29c
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/python/egglog/egraph.py b/python/egglog/egraph.py
@@ -1372,6 +1372,7 @@ def saturate(
         *,
         expr: Expr | None = None,
         max: int = 1000,
+        visualize: bool = True,
         **kwargs: Unpack[GraphvizKwargs],
     ) -> None:
         """
@@ -1387,12 +1388,15 @@ def to_json() -> str:
                 print(self.extract(expr), "\n")
             return self._serialize(**kwargs).to_json()
 
-        egraphs = [to_json()]
+        if visualize:
+            egraphs = [to_json()]
         i = 0
         while self.run(schedule or 1).updated and i < max:
             i += 1
-            egraphs.append(to_json())
-        VisualizerWidget(egraphs=egraphs).display_or_open()
+            if visualize:
+                egraphs.append(to_json())
+        if visualize:
+            VisualizerWidget(egraphs=egraphs).display_or_open()
 
     @classmethod
     def current(cls) -> EGraph: