Zip download fix (supersedes #7)

src/posteriordb/dataset.py

@@ -1,5 +1,4 @@
 import json
-import os
 import tempfile
 from typing import Union
 from zipfile import ZipFile

src/posteriordb/model.py

@@ -1,4 +1,3 @@
-import os
 from typing import Union
 
 from .posterior_database import PosteriorDatabase

src/posteriordb/posterior.py

@@ -38,6 +38,11 @@ def reference_draws_file_path(self):
 
     def reference_draws(self):
         reference_name = self.reference_draws_info()["name"]
-        with ZipFile(self.reference_draws_file_path() + ".zip", "r") as z:
-            with z.open(reference_name + ".json", "r") as f:
-                return json.load(f)
+        file_path = self.reference_draws_file_path()
+        if file_path.exists():
+            with ZipFile(file_path, "r") as z:
+                with z.open(reference_name + ".json", "r") as f:
+                    reference_draws =  json.load(f)
+        else:
+            raise ValueError(f"Missing reference_draws: {file_path}")
+        return reference_draws

src/posteriordb/posterior_database.py

@@ -84,8 +84,8 @@ def get_reference_draws_path(self, name: str):
         reference_root = os.path.join(self.path, "reference_posteriors/draws/draws")
         reference_name = self.get_posterior_info(name)["reference_posterior_name"]
         assert reference_name is not None
-        file_path = os.path.join(reference_root, reference_name + ".json")
-        return file_path
+        file_path = os.path.join(reference_root, reference_name + ".json.zip")
+        return Path(file_path)
 
     def get_reference_draws_info(self, name: str):
         reference_root = os.path.join(self.path, "reference_posteriors/draws/info")

src/posteriordb/posterior_database_github.py

@@ -1,7 +1,6 @@
 import hashlib
 import json
 import os
-import tempfile
 from pathlib import Path
 import requests
 
@@ -131,6 +130,24 @@ def load_info(path, assertion_function, metadata, overwrite=False):
     return info
 
 
+def load_zip_file(path, metadata, overwrite=False):
+    if not path.exists():
+        if metadata is not None:
+            download_file(metadata["download_url"], path)
+        else:
+            raise TypeError("File was not found in GitHub")
+    elif overwrite:
+        if metadata is not None:
+            download_file(metadata["download_url"], path, sha=metadata["sha"])
+    return path
+
+
+def load_draws(path, assertion_function, metadata, overwrite=False):
+    path = load_zip_file(path, metadata, overwrite=overwrite)
+    assertion_function(path)
+    return path
+
+
 def filenames_in_dir_no_extension(directory, gh_directory, extension):
     path = Path(directory)
 
@@ -266,8 +283,13 @@ def get_reference_draws_path(self, name: str):
         reference_root = self.path / "reference_posteriors" / "draws" / "draws"
         reference_name = self.get_posterior_info(name).get("reference_posterior_name")
         assert reference_name is not None
-        path = reference_root / (reference_name + ".json")
-        return path
+        path = reference_root / (reference_name + ".json.zip")
+        return load_draws(
+            path,
+            temporary_no_assertions,
+            self._links.get(path, None),
+            overwrite=self.overwrite,
+        )
 
     def get_reference_draws_info(self, name: str):
         reference_root = self.path / "reference_posteriors" / "draws" / "info"