move abstract b2share to abstract file

mlair/reference_data_handler/abstract_reference_data_handler.py

@@ -2,8 +2,13 @@ __author__ = "Felix Kleinert"
 __date__ = "2021-01-29"
 
 import os
+import sys
 from abc import ABC
 
+import wget
+
+from mlair.configuration import check_path_and_create
+
 
 class AbstractReferenceModel(ABC):
     """
@@ -30,3 +35,30 @@ class AbstractReferenceModel(ABC):
             res = False
         return res
 
+
+class AbstractReferenceb2share(AbstractReferenceModel):
+    """
+    Abstract class for reference models located on b2share (eudat or fz-juelich)
+    See also https://github.com/EUDAT-Training/B2SHARE-Training/blob/master/api/01_Retrieve_existing_record.md
+
+    """
+    def __init__(self, b2share_hosturl: str, b2share_bucket: str, b2share_key: str):
+        super().__init__()
+        self.b2share_hosturl = b2share_hosturl
+        self.b2share_bucket = b2share_bucket
+        self.b2share_key = b2share_key
+
+    @property
+    def b2share_url(self):
+        return f"{self.b2share_hosturl}/api/files/{self.b2share_bucket}"
+
+    def bar_custom(self, current, total, width=80):
+        progress_message = f"Downloading {self.b2share_key}: {round(current / total * 100)}% [{current} / {total}] bytes"
+        sys.stdout.write("\r" + progress_message)
+        sys.stdout.flush()
+
+    def download_from_b2share(self, tmp_download_path: str):
+        check_path_and_create(tmp_download_path)
+        wget.download(f"{self.b2share_url}/{self.b2share_key}",
+                      out=f"{tmp_download_path}{self.b2share_key}",
+                      bar=self.bar_custom)
\ No newline at end of file

mlair/reference_data_handler/IntelliO3v1Reference.py → mlair/reference_data_handler/intellio3_v1_reference.py

@@ -8,40 +8,10 @@ __date__ = "2021-01-29"
 
 import os
 import xarray as xr
-import wget
-import sys
 import shutil
 
 from mlair.configuration.path_config import check_path_and_create
-from mlair.reference_data_handler.abstract_reference_data_handler import AbstractReferenceModel
-
-
-class AbstractReferenceb2share(AbstractReferenceModel):
-    """
-    Abstract class for reference models located on b2share (eudat or fz-juelich)
-    See also https://github.com/EUDAT-Training/B2SHARE-Training/blob/master/api/01_Retrieve_existing_record.md
-
-    """
-    def __init__(self, b2share_hosturl: str, b2share_bucket: str, b2share_key: str):
-        super().__init__()
-        self.b2share_hosturl = b2share_hosturl
-        self.b2share_bucket = b2share_bucket
-        self.b2share_key = b2share_key
-
-    @property
-    def b2share_url(self):
-        return f"{self.b2share_hosturl}/api/files/{self.b2share_bucket}"
-
-    def bar_custom(self, current, total, width=80):
-        progress_message = f"Downloading {self.b2share_key}: {round(current / total * 100)}% [{current} / {total}] bytes"
-        sys.stdout.write("\r" + progress_message)
-        sys.stdout.flush()
-
-    def download_from_b2share(self, tmp_download_path: str):
-        check_path_and_create(tmp_download_path)
-        wget.download(f"{self.b2share_url}/{self.b2share_key}",
-                      out=f"{tmp_download_path}{self.b2share_key}",
-                      bar=self.bar_custom)
+from mlair.reference_data_handler.abstract_reference_data_handler import AbstractReferenceb2share
 
 
 class IntelliO3Reference(AbstractReferenceb2share):
@@ -107,7 +77,7 @@ class IntelliO3Reference(AbstractReferenceb2share):
             data.coords['type'] = (self.ref_name)
             data.to_netcdf(f"{self.ref_store_path}{infile}")
 
-    def make_reference_available_locally(self):
+    def make_reference_available_locally(self, remove_tmp_dir: bool = True):
         """
 
         :return:
@@ -118,7 +88,8 @@ class IntelliO3Reference(AbstractReferenceb2share):
                 self.download_from_b2share(tmp_download_path=self.tmp_extract_path)
             self.untar_forecasts()
             self.read_and_drop()
-            shutil.rmtree(self.tmp_extract_path)
+            if remove_tmp_dir:
+                shutil.rmtree(self.tmp_extract_path)
 
 
 if __name__ == '__main__':