From 96a10341271fb6377cd6c3e1cc0ff57ed7782c70 Mon Sep 17 00:00:00 2001
From: Shreya Shankar <ss.shankar505@gmail.com>
Date: Wed, 25 Sep 2024 21:49:31 -0700
Subject: [PATCH] Add a save config to the Python API

---
 docetl/api.py | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/docetl/api.py b/docetl/api.py
index ccde6f64..4647572f 100644
--- a/docetl/api.py
+++ b/docetl/api.py
@@ -47,9 +47,13 @@
 import os
 from typing import List, Optional, Dict, Any, Union
 
+import yaml
+
 from docetl.builder import Optimizer
 from docetl.runner import DSLRunner
 
+from rich import print
+
 
 @dataclass
 class Dataset:
@@ -274,6 +278,22 @@ def run(self, max_threads: Optional[int] = None) -> float:
         result = runner.run()
         return result
 
+    def to_yaml(self, path: str) -> None:
+        """
+        Convert the Pipeline object to a YAML string and save it to a file.
+
+        Args:
+            path (str): Path to save the YAML file.
+
+        Returns:
+            None
+        """
+        config = self._to_dict()
+        with open(path, "w") as f:
+            yaml.safe_dump(config, f)
+
+        print(f"[green]Pipeline saved to {path}[/green]")
+
     def _to_dict(self) -> Dict[str, Any]:
         """
         Convert the Pipeline object to a dictionary representation.