Updated retry logic so that it worked with storing release statuses a…

…s enums
v7labs · Jul 16, 2024 · 241c527 · 241c527
1 parent 2854a55
commit 241c527
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/darwin/cli_functions.py b/darwin/cli_functions.py
@@ -468,7 +468,7 @@ def pull_dataset(
         print_new_version_info(client)
     except NotFound:
         _error(
-            f"Version '{dataset.identifier}:{version}' does not exist "
+            f"Version '{dataset.identifier}:{version}' does not exist. "
             f"Use 'darwin dataset releases' to list all available versions."
         )
     except UnsupportedExportFormat as uef:

diff --git a/darwin/dataset/release.py b/darwin/dataset/release.py
@@ -94,7 +94,7 @@ def __init__(
         self.team_slug = team_slug
         self.version = version
         self.name = name
-        self.status = status.value
+        self.status = ReleaseStatus(status)
         self.url = url
         self.export_date = export_date
         self.image_count = image_count

diff --git a/darwin/dataset/remote_dataset.py b/darwin/dataset/remote_dataset.py
@@ -274,16 +274,16 @@ def pull(
         if release.format != "json" and release.format != "darwin_json_2":
             raise UnsupportedExportFormat(release.format)
 
-        if release.status == "pending":
+        if release.status.value == "pending":
             if retry:
-                while release.status == "pending" and retry_timeout > 0:
+                while release.status.value == "pending" and retry_timeout > 0:
                     console.print(
                         f"Release '{release.name}' for dataset '{self.name}' is still processing. Retrying in {retry_interval} seconds... {retry_timeout} seconds left before timeout."
                     )
                     time.sleep(retry_interval)
                     retry_timeout -= retry_interval
                     release = self.get_release(release.name, include_unavailable=retry)
-                if release.status == "pending":
+                if release.status.value == "pending":
                     raise ValueError(
                         f"Release {release.name} for dataset '{self.name}' is still processing. Please try again later."
                     )