feat: make gnuplot dataset writing more flexible (#886)

WilliamHPNielsen · jenshnielsen · commit 4c7ed38f689a · 2017-11-28T14:59:28.000+01:00
* feat: make gnuplot dataset writing more flexible

Allow users to save multiple files with a custom filename to a single
folder containing no snapshots

* fix: remove debug print

* fix: don't pass filename kwarg to hdf5 formatters
diff --git a/qcodes/data/data_set.py b/qcodes/data/data_set.py
@@ -469,7 +469,7 @@ def read_metadata(self):
             return
         self.formatter.read_metadata(self)
 
-    def write(self, write_metadata=False, only_complete=True):
+    def write(self, write_metadata=False, only_complete=True, filename=None):
         """
         Writes updates to the DataSet to storage.
         N.B. it is recommended to call data_set.finalize() when a DataSet is
@@ -480,15 +480,26 @@ def write(self, write_metadata=False, only_complete=True):
             only_complete (bool): passed on to the match_save_range inside
                 self.formatter.write. Used to ensure that all new data gets
                 saved even when some columns are strange.
+            filename (Optional[str]): The filename (minus extension) to use.
+                The file gets saved in the usual location.
         """
         if self.location is False:
             return
 
-        self.formatter.write(self,
-                             self.io,
-                             self.location,
-                             write_metadata=write_metadata,
-                             only_complete=only_complete)
+        # Only the gnuplot formatter has a "filename" kwarg
+        if isinstance(self.formatter, GNUPlotFormat):
+            self.formatter.write(self,
+                                 self.io,
+                                 self.location,
+                                 write_metadata=write_metadata,
+                                 only_complete=only_complete,
+                                 filename=filename)
+        else:
+            self.formatter.write(self,
+                                 self.io,
+                                 self.location,
+                                 write_metadata=write_metadata,
+                                 only_complete=only_complete)
 
     def write_copy(self, path=None, io_manager=None, location=None):
         """
@@ -562,21 +573,28 @@ def save_metadata(self):
             self.snapshot()
             self.formatter.write_metadata(self, self.io, self.location)
 
-    def finalize(self):
+    def finalize(self, filename=None, write_metadata=True):
         """
         Mark the DataSet complete and write any remaining modifications.
 
         Also closes the data file(s), if the ``Formatter`` we're using
         supports that.
+
+        Args:
+            filename (Optional[str]): The file name (minus extension) to
+                write to. The location of the file is the usual one.
+            write_metadata (bool): Whether to save a snapshot. For e.g. dumping
+                raw data inside a loop, a snapshot is not wanted.
         """
         log.debug('Finalising the DataSet. Writing.')
         # write all new data, not only (to?) complete columns
-        self.write(only_complete=False)
+        self.write(only_complete=False, filename=filename)
 
         if hasattr(self.formatter, 'close_file'):
             self.formatter.close_file(self)
 
-        self.save_metadata()
+        if write_metadata:
+            self.save_metadata()
 
     def snapshot(self, update=False):
         """JSON state of the DataSet."""
diff --git a/qcodes/data/gnuplot_format.py b/qcodes/data/gnuplot_format.py
@@ -244,7 +244,7 @@ def _get_labels(self, labelstr):
             return [l.replace('\\"', '"').replace('\\\\', '\\') for l in parts]
 
     def write(self, data_set, io_manager, location, force_write=False,
-              write_metadata=True, only_complete=True):
+              write_metadata=True, only_complete=True, filename=None):
         """
         Write updates in this DataSet to storage.
 
@@ -259,6 +259,9 @@ def write(self, data_set, io_manager, location, force_write=False,
                 or only complete rows? Is used to make sure that everything
                 gets written when the DataSet is finalised, even if some
                 dataarrays are strange (like, full of nans)
+            filename (Optional[str]): Filename to save to. Will override
+                the usual naming scheme and possibly overwrite files, so
+                use with care. The file will be saved in the normal location.
         """
         arrays = data_set.arrays
 
@@ -271,7 +274,11 @@ def write(self, data_set, io_manager, location, force_write=False,
         for group in groups:
             log.debug('Attempting to write the following '
                       'group: {}'.format(group))
-            fn = io_manager.join(location, group.name + self.extension)
+
+            if filename:
+                fn = io_manager.join(location, filename + self.extension)
+            else:
+                fn = io_manager.join(location, group.name + self.extension)
 
             written_files.add(fn)
 
