Merge pull request #153 from metno/152-delete-temporary-files-in-finalize

152 delete temporary files in finalize

Author: havardhhaugen

bris/outputs/grib.py

@@ -25,6 +25,7 @@ def __init__(
         proj4_str=None,
         domain_name=None,
         extra_variables=None,
+        remove_intermediate: bool = True,
     ):
         """
         Args:
@@ -43,6 +44,7 @@ def __init__(
         self.intermediate = None
         if self.pm.num_members > 1:
             self.intermediate = Intermediate(predict_metadata, workdir)
+        self.remove_intermediate = remove_intermediate
 
         self.variable_list = VariableList(self.extract_variables)
 
@@ -303,6 +305,9 @@ def finalize(self):
                 filename = self.get_filename(forecast_reference_time)
                 self.write(filename, forecast_reference_time, pred)
 
+            if self.remove_intermediate:
+                self.intermediate.cleanup()
+
     def get_lower(self, array):
         m = np.min(array)
         return np.floor(m / self.interp_res) * self.interp_res

bris/outputs/intermediate.py

@@ -95,8 +95,19 @@ def num_members(self):
     def get_filenames(self):
         return glob.glob(f"{self.workdir}/*_*.npy")
 
-    def finalize(self):
-        # clean up files
+    def cleanup(self):
+        """Removes up all intermediate files and removes the workdir. Called in finalize of the main output."""
+
         for _filename in self.get_filenames():
-            # delete file
-            pass
+            try:
+                os.remove(_filename)
+            except OSError as e:
+                print(f"Error during cleanup of {_filename}: {e}")
+
+        try:
+            os.rmdir(self.workdir)
+        except OSError as e:
+            print(f"Error removing workdir {self.workdir}: {e}")
+
+    def finalize(self):
+        pass

bris/outputs/netcdf.py

@@ -43,6 +43,7 @@ def __init__(
         mask_file=None,
         mask_field=None,
         global_attributes=None,
+        remove_intermediate=True,
     ):
         """
         Args:
@@ -61,8 +62,12 @@ def __init__(
 
         self.intermediate = None
         if self.pm.num_members > 1:
-            self.intermediate = Intermediate(predict_metadata, workdir, extra_variables)
-
+            self.intermediate = Intermediate(
+                predict_metadata,
+                workdir,
+                extra_variables,
+            )
+        self.remove_intermediate = remove_intermediate
         self.variable_list = VariableList(self.extract_variables)
 
         # Conventions specify the names of variables in the output
@@ -440,6 +445,9 @@ def finalize(self):
                 ]
                 self.write(filename, lead_times, pred)
 
+            if self.remove_intermediate:
+                self.intermediate.cleanup()
+
     def get_lower(self, array):
         m = np.min(array)
         return np.floor(m / self.interp_res) * self.interp_res

bris/outputs/verif.py

@@ -33,6 +33,7 @@ def __init__(
         fair_threshold: bool = True,
         fair_quantile: bool = True,
         fair_crps: bool = True,
+        remove_intermediate: bool = True,
     ) -> None:
         """
         Args:
@@ -125,6 +126,7 @@ def __init__(
             predict_metadata.num_members,
         )
         self.intermediate = Intermediate(intermediate_pm, workdir)
+        self.remove_intermediate = remove_intermediate
 
     def _add_forecast(self, times: list, ensemble_member: int, pred: np.array) -> None:
         """Add forecasts to this object. Will be written when .write() is called
@@ -387,6 +389,9 @@ def finalize(self) -> None:
         utils.create_directory(self.filename)
         self.ds.to_netcdf(self.filename, mode="w", engine="netcdf4")
 
+        if self.remove_intermediate:
+            self.intermediate.cleanup()
+
     def compute_consensus(self, pred) -> np.ndarray:
         assert len(pred.shape) == 3, pred.shape