'Remove step' feature for local workspaces

tango/workspace.py

@@ -20,6 +20,8 @@
 from urllib.parse import ParseResult, urlparse
 
 import pytz
+from sqlitedict import SqliteDict
+import shutil
 
 from .common import Registrable
 from .common.from_params import FromParams
@@ -419,6 +421,29 @@ def step_result(self, step_name: str) -> Any:
                 return self.step_cache[run.steps[step_name]]
         raise KeyError(f"No step named '{step_name}' found in previous runs")
 
+
+    def remove_step(self, step_name: str) -> Any:
+        """
+        Get Step name (Unique ID) from the user and remove the step information from cache
+        :raises KeyError: If no step with the unique name found in the cache dir
+        """
+        # get path to dir dynamically
+        sqlite_path = self.dir / "stepinfo.sqlite"
+        with SqliteDict(sqlite_path) as d:
+            try:
+                step_location = d[step_name].result_location
+                # remove step info from the sqlite dict
+                del d[step_name]
+                d.commit()
+                # remove cache directory
+                try:
+                    shutil.rmtree(step_location)
+                except OSError:
+                    raise OSError('Step Cache folder not found')
+                return('Step deleted')
+            except KeyError:
+                raise KeyError(f"No step named '{step_name}' found")
+
     def capture_logs_for_run(self, name: str) -> ContextManager[None]:
         """
         Should return a context manager that can be used to capture the logs for a run.