Refactored code for saving files

filedialogs.get_file_name_save now uses filedialogs.get_filename_format

Orange/widgets/data/owfile.py

@@ -18,7 +18,7 @@
 from Orange.widgets.utils.domaineditor import DomainEditor
 from Orange.widgets.utils.itemmodels import PyListModel
 from Orange.widgets.utils.filedialogs import RecentPathsWComboMixin, \
-    get_file_name_open
+    get_filename_format
 from Orange.widgets.widget import Output
 
 # Backward compatibility: class RecentPath used to be defined in this module,
@@ -269,7 +269,7 @@ def browse_file(self, in_demos=False):
             start_file = self.last_path() or os.path.expanduser("~/")
 
         readers = FileFormat.get_file_readers()
-        filename, reader, _ = get_file_name_open(start_file, None, readers)
+        filename, reader, _ = get_filename_format(start_file, None, readers, add_all=True)
         if not filename:
             return
         self.add_path(filename)

Orange/widgets/data/tests/test_owsave.py

@@ -0,0 +1,71 @@
+# Test methods with long descriptive names can omit docstrings
+# pylint: disable=missing-docstring
+from unittest.mock import patch
+import os
+
+from Orange.data import Table
+from Orange.data.io import TabReader, PickleReader
+from Orange.tests import named_file
+from Orange.widgets.tests.base import WidgetTest
+from Orange.widgets.utils.filedialogs import format_filter, fix_extension
+from Orange.widgets.data.owsave import OWSave
+
+
+class TestOWSave(WidgetTest):
+
+    def setUp(self):
+        self.widget = self.create_widget(OWSave)  # type: OWSave
+
+    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:
+                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_fix_extension(self):
+        self.send_signal(self.widget.Inputs.data, Table("iris"))
+
+        # need to add vars
+        def mock_choice(ret):
+            f = lambda a, b, c, d: ret
+            for a, v in vars(fix_extension).items():
+                setattr(f, a, v)
+            return f
+
+        fix_change_ext = mock_choice(fix_extension.CHANGE_EXT)
+        fix_change_format = mock_choice(fix_extension.CHANGE_FORMAT)
+        fix_cancel = mock_choice(fix_extension.CANCEL)
+
+        def pickle_file_tab_filter(a, b, c, d, e):
+            return filename.replace("tab", "pickle"), format_filter(TabReader)
+
+        def tab_file_pickle_filter(a, b, c, d, e):
+            return filename, format_filter(PickleReader)
+
+        counter = [0]
+
+        def tab_file_change_filters(a, b, c, d, e):
+            counter[0] += 1
+            if counter[0] == 1:
+                return filename, format_filter(PickleReader)
+            else:
+                return filename, format_filter(TabReader)
+
+        with named_file("", suffix=".tab") as filename:
+
+            for file_choice, fix in [(pickle_file_tab_filter, fix_change_ext),
+                                     (tab_file_pickle_filter, fix_change_format),
+                                     (tab_file_change_filters, fix_cancel)]:
+
+                os.remove(filename)
+
+                with patch("AnyQt.QtWidgets.QFileDialog.getSaveFileName", file_choice),\
+                     patch("Orange.widgets.utils.filedialogs.fix_extension", fix):
+                    self.widget.save_file_as()
+
+                self.assertEqual(len(Table(filename)), 150)

Orange/widgets/utils/filedialogs.py

@@ -55,81 +55,84 @@ def dialog_formats():
 
 
 def get_file_name(start_dir, start_filter, file_formats):
-    """
-    Get filename for the given possible file formats
+    return get_file_name_save(start_dir, start_filter,
+                              sorted(set(file_formats.values()), key=lambda x: x.PRIORITY))
+
 
+def get_file_name_save(start_dir, start_filter, file_formats):
+    """
     The function uses the standard save file dialog with filters from the
     given file formats. Extension is added automatically, if missing. If the
     user enters file extension that does not match the file format, (s)he is
     given a dialog to decide whether to fix the extension or the format.
 
-    Function also returns the writer and filter to cover the case where the
-    same extension appears in multiple filters. Although `file_format` is a
-    dictionary that associates its extension with one writer, writers can
-    still have other extensions that are allowed.
-
     Args:
         start_dir (str): initial directory, optionally including the filename
         start_filter (str): initial filter
-        file_formats (dict {extension: Orange.data.io.FileFormat}): file formats
+        file_formats (a list of Orange.data.io.FileFormat): file formats
     Returns:
-        (filename, filter, writer), or `(None, None, None)` on cancel
+        (filename, file_format, filter), or `(None, None, None)` on cancel
     """
-    writers = sorted(set(file_formats.values()), key=lambda w: w.PRIORITY)
-    filters = [format_filter(w) for w in writers]
-    if start_filter not in filters:
-        start_filter = filters[0]
-
+    file_formats = list(file_formats)
     while True:
-        filename, filter = QFileDialog.getSaveFileName(
-            None, 'Save As...', start_dir, ';;'.join(filters), start_filter)
+        dialog = QFileDialog.getSaveFileName
+        filename, format, filter = \
+            get_filename_format(start_dir, start_filter, file_formats,
+                                add_all=False, title="Save as...", dialog=dialog)
         if not filename:
             return None, None, None
 
-        writer = writers[filters.index(filter)]
         base, ext = os.path.splitext(filename)
         if not ext:
-            filename += writer.EXTENSIONS[0]
-        elif ext not in writer.EXTENSIONS:
-            format = writer.DESCRIPTION
-            suggested_ext = writer.EXTENSIONS[0]
-            suggested_format = \
-                ext in file_formats and file_formats[ext].DESCRIPTION
-            res = fix_extension(ext, format, suggested_ext, suggested_format)
+            filename += format.EXTENSIONS[0]
+        elif ext not in format.EXTENSIONS:
+            suggested_ext = format.EXTENSIONS[0]
+            suggested_format = False
+            for f in file_formats:  # find the first format
+                if ext in f.EXTENSIONS:
+                    suggested_format = f
+                    break
+            res = fix_extension(ext, format.DESCRIPTION, suggested_ext,
+                                suggested_format.DESCRIPTION if suggested_format else False)
             if res == fix_extension.CANCEL:
                 continue
             if res == fix_extension.CHANGE_EXT:
                 filename = base + suggested_ext
             elif res == fix_extension.CHANGE_FORMAT:
-                writer = file_formats[ext]
-                filter = format_filter(writer)
-        return filename, writer, filter
+                format = suggested_format
+                filter = format_filter(format)
+        return filename, format, filter
 
 
-def get_file_name_open(start_dir, start_filter, file_formats, title="Open...",
-                       dialog=None):
+def get_filename_format(start_dir, start_filter, file_formats,
+                        add_all=False, title="Open...", dialog=None):
     """
     Open file dialog with file formats.
 
+    Function also returns the format and filter to cover the case where the
+    same extension appears in multiple filters.
+
     Args:
         start_dir (str): initial directory, optionally including the filename
         start_filter (str): initial filter
         file_formats (a list of Orange.data.io.FileFormat): file formats
+        add_all (bool): add a filter for all supported extensions
         title (str): title of the dialog
         dialog: a function that creates a QT dialog
     Returns:
-        (filename, filter, file_format), or `(None, None, None)` on cancel
+        (filename, file_format, filter), or `(None, None, None)` on cancel
     """
     file_formats = list(file_formats)
     filters = [format_filter(f) for f in file_formats]
 
     # add all readable files option
-    all_extensions = set()
-    for f in file_formats:
-        all_extensions.update(f.EXTENSIONS)
-    file_formats.insert(0, None)
-    filters.insert(0, "All readable files (*{})".format(
-        ' *'.join(sorted(all_extensions))))
+    if add_all:
+        all_extensions = set()
+        for f in file_formats:
+            all_extensions.update(f.EXTENSIONS)
+        file_formats.insert(0, None)
+        filters.insert(0, "All readable files (*{})".format(
+            ' *'.join(sorted(all_extensions))))
 
     if start_filter not in filters:
         start_filter = filters[0]