flora using folders

Author: emptymalei

dataherb/flora.py

@@ -1,14 +1,15 @@
 import json
+import shutil
 import sys
 from pathlib import Path
 
 from loguru import logger
 
 from dataherb.core.base import Herb
-from dataherb.core.search import search_by_ids_in_flora as _search_by_ids_in_flora
-from dataherb.core.search import (
-    search_by_keywords_in_flora as _search_by_keywords_in_flora,
-)
+from dataherb.core.search import \
+    search_by_ids_in_flora as _search_by_ids_in_flora
+from dataherb.core.search import \
+    search_by_keywords_in_flora as _search_by_keywords_in_flora
 from dataherb.fetch.remote import get_data_from_url as _get_data_from_url
 from dataherb.parse.model_json import MetaData
 
@@ -21,12 +22,15 @@ class Flora(object):
     DataHerb is the container of datasets.
     """
 
-    def __init__(self, flora):
+    def __init__(self, flora, is_aggregated=False):
         """
         :param flora: API of the DataHerb service, defaults to dataherb official or list of Herbs, defaults to everything from the API
         """
+        self.is_aggregated = is_aggregated
 
         if isinstance(flora, str):
+            if flora.endswith(".json"):
+                self.is_aggregated = True
             self.workdir = Path(flora).parent.parent
             self.flora_config = flora
             flora = self._get_flora(flora)
@@ -42,8 +46,18 @@ def _get_flora(self, flora_config):
         _get_flora fetch flora from the provided API.
         """
         if Path(flora_config).exists():
-            with open(flora_config, "r") as f:
-                flora = json.load(f)
+            if self.is_aggregated:
+                with open(flora_config, "r") as f:
+                    flora = json.load(f)
+            else:
+                # We do not have aggregated flora
+                # read dataherb.json from all the folders here
+                flora_folder = Path(flora_config)
+                herb_paths = [f for f in flora_folder.iterdir() if f.is_dir()]
+                flora = [
+                    json.load(open(f.joinpath("dataherb.json"), "r"))
+                    for f in herb_paths
+                ]
         else:
             # assuming the config is a url if the local file does not exist
             flora_request = _get_data_from_url(flora_config)
@@ -91,16 +105,18 @@ def remove(self, herb_id):
         """
         add add a herb to the flora
         """
-
         for id in [i.id for i in self.flora]:
             if id == herb_id:
                 logger.debug(f"found herb id = {herb_id}")
 
         self.flora = [h for h in self.flora if h.id != herb_id]
 
-        self.save()
+        if self.is_aggregated:
+            self.save()
+        else:
+            self.remove_herb_from_flora(herb_id)
 
-    def save(self, path=None):
+    def save(self, path=None, id=None):
         """save flora metadata to json file"""
 
         if path is None:
@@ -110,16 +126,51 @@ def save(self, path=None):
             f"type of a herb in flora: {type(self.flora[0])}\n{self.flora[0].metadata}"
         )
 
-        serialized_flora = []
-        for h in self.flora:
-            logger.debug(f"herb (type {type(h)}): {h}")
-            serialized_flora.append(h.metadata)
+        if self.is_aggregated:
+            serialized_flora = []
+            for h in self.flora:
+                logger.debug(f"herb (type {type(h)}): {h}")
+                serialized_flora.append(h.metadata)
+
+            with open(path, "w") as fp:
+                json.dump(
+                    serialized_flora, fp, sort_keys=True, indent=4, separators=(",", ": ")
+                )
+        else:
+            if not id:
+                raise Exception("dataherb id must be provided")
+
+    def save_herb_meta(self, id, path=None):
+        """Save a herb metadata to json file"""
+        if path is None:
+            path = self.workdir / f"{id}"
+
+        if not path.exists():
+            path.mkdir(parents=True)
 
-        with open(path, "w") as fp:
+        logger.debug(f"Will replace dataherb id {id}")
+        with open(path / "dataherb.json", "w") as fp:
             json.dump(
-                serialized_flora, fp, sort_keys=True, indent=4, separators=(",", ": ")
+                self.herb_meta(id), fp,
+                sort_keys=True, indent=4, separators=(",", ": ")
             )
 
+    def remove_herb_from_flora(self, id, path=None):
+        """Remove a herb metadata to json file"""
+        if path is None:
+            path = self.workdir / f"{id}"
+
+        if not path.exists():
+            logger.debug(f"dataherb {id} doesn't exist")
+            return
+        else:
+            try:
+                shutil.rmtree(path)
+            except OSError as e:
+                logger.error(
+                    f"Can not remove herb id {id}: {e.filename} - {e.strerror}."
+                )
+
     def search(self, keywords):
         """
         search finds the datasets that matches the keywords