Simplify dataset saving to disk

saving datasets can now be done one by one. If a writer is not provided,
it is guessed from the filename extension.

Signed-off-by: Martin Raspaud <martin.raspaud@smhi.se>

mpop/scene.py

@@ -469,7 +469,18 @@ def load_writer_config(self, config_files, **kwargs):
                 writer = writer_class(ppp_config_dir=self.ppp_config_dir, config_file=config_files, **kwargs)
                 return writer
 
-    def save_images(self, writer="geotiff", **kwargs):
+    def save_dataset(self, dataset_id, filename=None, writer=None, **kwargs):
+        """Save the *dataset_id* to file using *writer* (geotiff by default).
+        """
+        if writer is None:
+            writer = self.get_writer_by_ext(os.path.splitext(filename)[1], **kwargs)
+        else:
+            writer = self.get_writer(writer, **kwargs)
+        writer.save_dataset(self[dataset_id], filename=filename)
+
+    def save_datasets(self, writer="geotiff", **kwargs):
+        """Save all the datasets present in a scene to disk using *writer*.
+        """
         writer = self.get_writer(writer, **kwargs)
         for projectable in self.datasets.values():
             writer.save_dataset(projectable, **kwargs)
@@ -478,4 +489,10 @@ def get_writer(self, writer="geotiff", **kwargs):
         config_fn = writer + ".cfg" if "." not in writer else writer
         config_files = config_search_paths(os.path.join("writers", config_fn), self.ppp_config_dir)
         kwargs.setdefault("config_files", config_files)
-        return self.load_writer_config(**kwargs)
+        return self.load_writer_config(**kwargs)
+
+    def get_writer_by_ext(self, extension, **kwargs):
+        mapping = {".tiff": "geotiff",
+                   ".tif":  "geotiff",
+                   }
+        return self.get_writer(mapping.get(extension.lower(), "simple_image"), **kwargs)

mpop/writers/__init__.py

@@ -147,14 +147,14 @@ def get_filename(self, **kwargs):
             raise RuntimeError("No filename pattern or specific filename provided")
         return self.filename_parser.compose(kwargs)
 
-    def save_dataset(self, dataset, fill_value=None, **kwargs):
+    def save_dataset(self, dataset, filename=None, fill_value=None, **kwargs):
         """Saves the *dataset* to a given *filename*.
         """
         fill_value = fill_value if fill_value is not None else self.fill_value
         img = get_enhanced_image(dataset, self.enhancer, fill_value)
-        self.save_image(img, **kwargs)
+        self.save_image(img, filename=filename, **kwargs)
 
-    def save_image(self, img, *args, **kwargs):
+    def save_image(self, img, filename=None, **kwargs):
         raise NotImplementedError("Writer '%s' has not implemented image saving" % (self.name,))
 
 

mpop/writers/geotiff.py

@@ -100,7 +100,7 @@ def _gdal_write_datasets(self, dst_ds, datasets, opacity, fill_value):
             alpha = np.where(mask, 0, opacity).astype(chan.dtype)
             dst_ds.GetRasterBand(i + 2).WriteArray(alpha)
 
-    def save_image(self, img, floating_point=False, **kwargs):
+    def save_image(self, img, filename=None, floating_point=False, **kwargs):
         """Save the image to the given *filename* in geotiff_ format.
         `floating_point` allows the saving of
         'L' mode images in floating point format if set to True.
@@ -109,9 +109,7 @@ def save_image(self, img, floating_point=False, **kwargs):
         """
         raster = gdal.GetDriverByName("GTiff")
 
-        metadata = img.info
-
-        filename = kwargs.pop("filename", self.get_filename(**img.info))
+        filename = filename or self.get_filename(**img.info)
 
         # Update global GDAL options with these specific ones
         gdal_options = self.gdal_options.copy()

mpop/writers/simple_image.py

@@ -33,7 +33,8 @@ class PillowWriter(Writer):
     def __init__(self, **kwargs):
         Writer.__init__(self, default_config_filename="writers/simple_image.cfg", **kwargs)
 
-    def save_image(self, img, **kwargs):
-        filename = kwargs.pop("filename", self.get_filename(**img.info))
+    def save_image(self, img, filename=None, **kwargs):
+        filename = filename or self.get_filename(**img.info)
+
         LOG.debug("Saving to image: %s", filename)
         img.save(filename)