Merge pull request #3147 from golbog/owsave

[ENH] Selection of format and compression in save data widget

Orange/data/io.py

@@ -931,7 +931,7 @@ class TabReader(CSVReader):
 class PickleReader(FileFormat):
     """Reader for pickled Table objects"""
     EXTENSIONS = ('.pkl', '.pickle')
-    DESCRIPTION = 'Pickled Python object file'
+    DESCRIPTION = 'Pickled Orange data'
     SUPPORT_COMPRESSED = True
     SUPPORT_SPARSE_DATA = True
 

Orange/widgets/data/owsave.py

@@ -1,12 +1,29 @@
 import os.path
+import pathlib
+
+from AnyQt.QtWidgets import QFormLayout
+from AnyQt.QtCore import Qt
 
 from Orange.data.table import Table
+from Orange.data.io import Compression, FileFormat, TabReader, CSVReader, PickleReader
 from Orange.widgets import gui, widget
 from Orange.widgets.settings import Setting
-from Orange.data.io import FileFormat
 from Orange.widgets.utils import filedialogs
 from Orange.widgets.widget import Input
 
+FILE_TYPES = [
+    ("{} ({})".format(w.DESCRIPTION, w.EXTENSIONS[0]),
+     w.EXTENSIONS[0],
+     w.SUPPORT_SPARSE_DATA)
+    for w in (TabReader, CSVReader, PickleReader)
+]
+
+COMPRESSIONS = [
+    ("gzip ({})".format(Compression.GZIP), Compression.GZIP),
+    ("bzip2 ({})".format(Compression.BZIP2), Compression.BZIP2),
+    ("lzma ({})".format(Compression.XZ), Compression.XZ),
+]
+
 
 class OWSave(widget.OWWidget):
     name = "Save Data"
@@ -18,62 +35,130 @@ class OWSave(widget.OWWidget):
     class Inputs:
         data = Input("Data", Table)
 
+    class Error(widget.OWWidget.Error):
+        unsupported_extension = widget.Msg("Selected extension is not supported.")
+
     want_main_area = False
     resizing_enabled = False
 
     last_dir = Setting("")
-    last_filter = Setting("")
     auto_save = Setting(False)
-
-    @classmethod
-    def get_writers(cls, sparse):
-        return [f for f in FileFormat.formats
-                if getattr(f, 'write_file', None) and getattr(f, "EXTENSIONS", None)
-                and (not sparse or getattr(f, 'SUPPORT_SPARSE_DATA', False))]
+    filetype = Setting(FILE_TYPES[0][0])
+    compression = Setting(COMPRESSIONS[0][0])
+    compress = Setting(False)
 
     def __init__(self):
         super().__init__()
         self.data = None
         self.filename = ""
+        self.basename = ""
+        self.type_ext = ""
+        self.compress_ext = ""
         self.writer = None
 
+        form = QFormLayout(
+            labelAlignment=Qt.AlignLeft,
+            formAlignment=Qt.AlignLeft,
+            rowWrapPolicy=QFormLayout.WrapLongRows,
+            verticalSpacing=10,
+        )
+
+        box = gui.vBox(self.controlArea, "Format")
+
+        gui.comboBox(
+            box, self, "filetype",
+            callback=self._update_text,
+            items=[item for item, _, _ in FILE_TYPES],
+            sendSelectedValue=True,
+        )
+        form.addRow("File type", self.controls.filetype, )
+
+        gui.comboBox(
+            box, self, "compression",
+            callback=self._update_text,
+            items=[item for item, _ in COMPRESSIONS],
+            sendSelectedValue=True,
+        )
+        gui.checkBox(
+            box, self, "compress", label="Use compression",
+            callback=self._update_text,
+        )
+
+        form.addRow(self.controls.compress, self.controls.compression)
+
+        box.layout().addLayout(form)
+
         self.save = gui.auto_commit(
             self.controlArea, self, "auto_save", "Save", box=False,
             commit=self.save_file, callback=self.adjust_label,
-            disabled=True, addSpace=True)
-        self.saveAs = gui.button(
+            disabled=True, addSpace=True
+        )
+        self.save_as = gui.button(
             self.controlArea, self, "Save As...",
-            callback=self.save_file_as, disabled=True)
-        self.saveAs.setMinimumWidth(220)
+            callback=self.save_file_as, disabled=True
+        )
+        self.save_as.setMinimumWidth(220)
         self.adjustSize()
 
+    def get_writer_selected(self):
+        writer = FileFormat.get_reader(self.type_ext)
+
+        ext = self.type_ext + self.compress_ext
+        if ext not in writer.EXTENSIONS:
+            self.Error.unsupported_extension()
+            return None
+        writer.EXTENSIONS = [ext]
+        return writer
+
+    @classmethod
+    def remove_extensions(cls, filename):
+        if not filename:
+            return None
+        for ext in pathlib.PurePosixPath(filename).suffixes:
+            filename = filename.replace(ext, '')
+        return filename
+
     def adjust_label(self):
         if self.filename:
-            filename = os.path.split(self.filename)[1]
-            text = ["Save as '{}'", "Auto save as '{}'"][self.auto_save]
-            self.save.button.setText(text.format(filename))
+            text = "Auto save as '{}'" if self.auto_save else "Save as '{}'"
+            self.save.button.setText(
+                text.format(self.basename + self.type_ext + self.compress_ext))
 
     @Inputs.data
     def dataset(self, data):
         self.data = data
         self.save.setDisabled(data is None)
-        self.saveAs.setDisabled(data is None)
+        self.save_as.setDisabled(data is None)
         if data is not None:
             self.save_file()
 
+        self.controls.filetype.clear()
+        if self.data.is_sparse():
+            self.controls.filetype.insertItems(0, [item for item, _, supports_sparse in FILE_TYPES
+                                                   if supports_sparse])
+        else:
+            self.controls.filetype.insertItems(0, [item for item, _, _ in FILE_TYPES])
+
     def save_file_as(self):
-        file_name = self.filename or \
-            os.path.join(self.last_dir or os.path.expanduser("~"),
-                         getattr(self.data, 'name', ''))
-        filename, writer, filter = filedialogs.open_filename_dialog_save(
-            file_name, self.last_filter, self.get_writers(self.data.is_sparse()))
+        file_name = self.remove_extensions(self.filename) or os.path.join(
+            self.last_dir or os.path.expanduser("~"),
+            getattr(self.data, 'name', ''))
+        self.update_extension()
+        writer = self.get_writer_selected()
+        if not writer:
+            return
+
+        filename, writer, _ = filedialogs.open_filename_dialog_save(
+            file_name, '', [writer],
+        )
         if not filename:
             return
+
         self.filename = filename
         self.writer = writer
-        self.unconditional_save_file()
         self.last_dir = os.path.split(self.filename)[0]
-        self.last_filter = filter
+        self.basename = os.path.basename(self.remove_extensions(filename))
+        self.unconditional_save_file()
         self.adjust_label()
 
     def save_file(self):
@@ -83,18 +168,30 @@ def save_file(self):
             self.save_file_as()
         else:
             try:
-                self.writer.write(self.filename, self.data)
-            except Exception as errValue:
-                self.error(str(errValue))
+                self.writer.write(os.path.join(self.last_dir,
+                                               self.basename + self.type_ext + self.compress_ext),
+                                  self.data)
+            except Exception as err_value:
+                self.error(str(err_value))
             else:
                 self.error()
 
+    def update_extension(self):
+        self.type_ext = [ext for name, ext, _ in FILE_TYPES if name == self.filetype][0]
+        self.compress_ext = dict(COMPRESSIONS)[self.compression] if self.compress else ''
+
+    def _update_text(self):
+        self.update_extension()
+        self.adjust_label()
+
 
 if __name__ == "__main__":
     import sys
     from AnyQt.QtWidgets import QApplication
+
     a = QApplication(sys.argv)
     table = Table("iris")
+
     ow = OWSave()
     ow.show()
     ow.dataset(table)

Orange/widgets/data/tests/test_owsave.py

@@ -1,14 +1,28 @@
 # Test methods with long descriptive names can omit docstrings
 # pylint: disable=missing-docstring
 from unittest.mock import patch
+import itertools
 
 from Orange.data import Table
-from Orange.data.io import TabReader, PickleReader, FileFormat
+from Orange.data.io import Compression, FileFormat, TabReader, CSVReader, PickleReader
 from Orange.tests import named_file
 from Orange.widgets.tests.base import WidgetTest
-from Orange.widgets.utils.filedialogs import format_filter, fix_extension, open_filename_dialog_save
+from Orange.widgets.utils.filedialogs import format_filter
 from Orange.widgets.data.owsave import OWSave
 
+FILE_TYPES = [
+    ("{} ({})".format(w.DESCRIPTION, w.EXTENSIONS[0]),
+     w.EXTENSIONS[0],
+     w.SUPPORT_SPARSE_DATA)
+    for w in (TabReader, CSVReader, PickleReader)
+]
+
+COMPRESSIONS = [
+    ("gzip ({})".format(Compression.GZIP), Compression.GZIP),
+    ("bzip2 ({})".format(Compression.BZIP2), Compression.BZIP2),
+    ("lzma ({})".format(Compression.XZ), Compression.XZ),
+]
+
 
 class AddedFormat(FileFormat):
     EXTENSIONS = ('.234',)
@@ -23,53 +37,48 @@ class TestOWSave(WidgetTest):
     def setUp(self):
         self.widget = self.create_widget(OWSave)  # type: OWSave
 
+    def test_writer(self):
+        compressions = [self.widget.controls.compression.itemText(i) for i in
+                        range(self.widget.controls.compression.count())]
+        types = [self.widget.controls.filetype.itemText(i)
+                 for i in range(self.widget.controls.filetype.count())]
+        for t, c, d in itertools.product(types, compressions, [True, False]):
+            self.widget.compression = c
+            self.widget.compress = d
+            self.widget.filetype = t
+            self.widget.update_extension()
+            self.assertEqual(len(self.widget.get_writer_selected().EXTENSIONS), 1)
+
     def test_ordinary_save(self):
         self.send_signal(self.widget.Inputs.data, Table("iris"))
 
-        for ext, writer in [('.tab', TabReader), ('.pickle', PickleReader)]:
-            with named_file("", suffix=ext) as filename:
+        for ext, suffix, _ in FILE_TYPES:
+            self.widget.filetype = ext
+            self.widget.update_extension()
+            writer = self.widget.get_writer_selected()
+            with named_file("", suffix=suffix) as filename:
                 def choose_file(a, b, c, d, e, fn=filename, w=writer):
                     return fn, format_filter(w)
+
                 with patch("AnyQt.QtWidgets.QFileDialog.getSaveFileName", choose_file):
                     self.widget.save_file_as()
                 self.assertEqual(len(Table(filename)), 150)
 
-    def test_filename_with_fix_extension(self):
-
-        def mock_fix_choice(ret):
-            f = lambda *x: ret
-            f.__dict__.update(fix_extension.__dict__)
-            return f
-
-        change_filter = iter([PickleReader, TabReader])
-
-        for file_choice, fix in [
-                [lambda *x: ("o.pickle", format_filter(TabReader)),
-                 mock_fix_choice(fix_extension.CHANGE_EXT)],
-                [lambda *x: ("o.tab", format_filter(PickleReader)),
-                 mock_fix_choice(fix_extension.CHANGE_FORMAT)],
-                [lambda *x: ("o.tab", format_filter(next(change_filter))),
-                 mock_fix_choice(fix_extension.CANCEL)]
-        ]:
-
-            with patch("AnyQt.QtWidgets.QFileDialog.getSaveFileName", file_choice),\
-                 patch("Orange.widgets.utils.filedialogs.fix_extension", fix):
-                saved_filename, format, filter = \
-                    open_filename_dialog_save(".", None, OWSave.get_writers(False))
-                self.assertEqual(saved_filename, "o.tab")
-                self.assertEqual(format, TabReader)
-                self.assertEqual(filter, format_filter(TabReader))
-
-    def test_added_format(self):
-        """Test that a format added after widget initialization is recognized"""
+    def test_compression(self):
         self.send_signal(self.widget.Inputs.data, Table("iris"))
-        called = False
-        with named_file("", suffix=".tab") as filename:
-            def test_format(sd, sf, ff, **kwargs):
-                nonlocal called
-                called = True
-                self.assertIn(AddedFormat, ff)
-                return filename, TabReader, ""
-            with patch("Orange.widgets.utils.filedialogs.open_filename_dialog", test_format):
-                self.widget.save_file_as()
-        self.assertTrue(called)
+
+        self.widget.compress = True
+        for type, compression in itertools.product([[x, ext] for x, ext, _ in FILE_TYPES],
+                                                   COMPRESSIONS):
+            self.widget.filetype = type[0]
+            self.widget.compression = compression[0]
+            self.widget.update_extension()
+            writer = self.widget.get_writer_selected()
+            with named_file("",
+                            suffix=type[1] + compression[1]) as filename:
+                def choose_file(a, b, c, d, e, fn=filename, w=writer):
+                    return fn, format_filter(w)
+
+                with patch("AnyQt.QtWidgets.QFileDialog.getSaveFileName", choose_file):
+                    self.widget.save_file_as()
+                self.assertEqual(len(Table(filename)), 150)