New import/export interface with dataset and databook `import_` and `export` methodsimport_export

and overloaded `import_set` and `import_book` functions.

4 files changed, 121 insertions, 31 deletions

diff --git a/tablib/core.py b/tablib/core.py
index 9db46c9..b1de323 100644
--- a/tablib/core.py
+++ b/tablib/core.py
@@ -153,6 +153,8 @@ class Dataset(object):
 
     """
 
+    _formats = {}
+
     def __init__(self, *args, **kwargs):
         self._data = list(Row(arg) for arg in args)
         self.__headers = None
@@ -254,11 +256,13 @@ class Dataset(object):
             try:
                 try:
                     setattr(cls, fmt.title, property(fmt.export_set, fmt.import_set))
+                    cls._formats[fmt.title] = (fmt.export_set, fmt.import_set)
                 except AttributeError:
                     setattr(cls, fmt.title, property(fmt.export_set))
+                    cls._formats[fmt.title] = (fmt.export_set, None)
 
             except AttributeError:
-                pass
+                cls._formats[fmt.title] = (None, None)
 
 
     def _validate(self, row=None, col=None, safety=False):
@@ -428,12 +432,34 @@ class Dataset(object):
             except TypeError:
                 return 0
 
+    def import_(self, format, in_stream, **kwargs):
+        """
+        Import `in_stream` to the :class:`Dataset` object using the `format`.
+
+        :param \*\*kwargs: (optional) custom configuration to the format `import_set`.
+        """
+        export_set, import_set = self._formats.get(format, (None, None))
+        if not import_set:
+            raise UnsupportedFormat
+
+        import_set(self, in_stream, **kwargs)
+
+    def export(self, format, **kwargs):
+        """
+        Export :class:`Dataset` object to `format`.
+
+        :param \*\*kwargs: (optional) custom configuration to the format `export_set`.
+        """
+        export_set, import_set = self._formats.get(format, (None, None))
+        if not export_set:
+            raise UnsupportedFormat
+
+        return export_set(self, **kwargs)
 
     # -------
     # Formats
     # -------
 
-
     @property
     def xls():
         """A Legacy Excel Spreadsheet representation of the :class:`Dataset` object, with :ref:`separators`. Cannot be set.
@@ -921,6 +947,8 @@ class Databook(object):
     """A book of :class:`Dataset` objects.
     """
 
+    _formats = {}
+
     def __init__(self, sets=None):
 
         if sets is None:
@@ -936,7 +964,6 @@ class Databook(object):
         except AttributeError:
             return '<databook object>'
 
-
     def wipe(self):
         """Removes all :class:`Dataset` objects from the :class:`Databook`."""
         self._datasets = []
@@ -949,11 +976,13 @@ class Databook(object):
             try:
                 try:
                     setattr(cls, fmt.title, property(fmt.export_book, fmt.import_book))
+                    cls._formats[fmt.title] = (fmt.export_book, fmt.import_book)
                 except AttributeError:
                     setattr(cls, fmt.title, property(fmt.export_book))
+                    cls._formats[fmt.title] = (fmt.export_book, None)
 
             except AttributeError:
-                pass
+                cls._formats[fmt.title] = (None, None)
 
     def sheets(self):
         return self._datasets
@@ -988,6 +1017,30 @@ class Databook(object):
         """The number of the :class:`Dataset` objects within :class:`Databook`."""
         return len(self._datasets)
 
+    def import_(self, format, in_stream, **kwargs):
+        """
+        Import `in_stream` to the :class:`Databook` object using the `format`.
+
+        :param \*\*kwargs: (optional) custom configuration to the format `import_book`.
+        """
+        export_book, import_book = self._formats.get(format, (None, None))
+        if not import_book:
+            raise UnsupportedFormat
+
+        import_book(self, in_stream, **kwargs)
+
+    def export(self, format, **kwargs):
+        """
+        Export :class:`Databook` object to `format`.
+
+        :param \*\*kwargs: (optional) custom configuration to the format `export_book`.
+        """
+        export_book, import_book = self._formats.get(format, (None, None))
+        if not export_book:
+            raise UnsupportedFormat
+
+        return export_book(self, **kwargs)
+
 
 def detect(stream):
     """Return (format, stream) of given stream."""
@@ -1000,32 +1053,43 @@ def detect(stream):
     return (None, stream)
 
 
-def import_set(stream):
+def import_set(stream, format=None, **kwargs):
     """Return dataset of given stream."""
-    (format, stream) = detect(stream)
+    if format:
+        format = get_formatter(format)
+    else:
+        format, stream = detect(stream)
 
+    data = Dataset()
     try:
-        data = Dataset()
-        format.import_set(data, stream)
+        format.import_set(data, stream, **kwargs)
         return data
-
     except AttributeError:
         return None
 
 
-def import_book(stream):
+def import_book(stream, format=None, **kwargs):
     """Return dataset of given stream."""
-    (format, stream) = detect(stream)
+    if format:
+        format = get_formatter(format)
+    else:
+        format, stream = detect(stream)
 
+    databook = Databook()
     try:
-        databook = Databook()
-        format.import_book(databook, stream)
+        format.import_book(databook, stream, **kwargs)
         return databook
-
     except AttributeError:
         return None
 
 
+def get_formatter(format):
+    for item in formats.available:
+        if item.title == format:
+            return item
+    raise UnsupportedFormat(format)
+
+
 class InvalidDatasetType(Exception):
     "Only Datasets can be added to a DataBook"
 
diff --git a/tablib/formats/_csv.py b/tablib/formats/_csv.py
index 7597395..7d29318 100644
--- a/tablib/formats/_csv.py
+++ b/tablib/formats/_csv.py
@@ -14,14 +14,14 @@ DEFAULT_ENCODING = 'utf-8'
 
 
 
-def export_set(dataset):
+def export_set(dataset, **kwargs):
     """Returns CSV representation of Dataset."""
     stream = StringIO()
 
-    if is_py3:
-        _csv = csv.writer(stream)
-    else:
-        _csv = csv.writer(stream, encoding=DEFAULT_ENCODING)
+    if not is_py3:
+        kwargs.setdefault('encoding', DEFAULT_ENCODING)
+
+    _csv = csv.writer(stream, **kwargs)
 
     for row in dataset._package(dicts=False):
         _csv.writerow(row)
@@ -29,15 +29,15 @@ def export_set(dataset):
     return stream.getvalue()
 
 
-def import_set(dset, in_stream, headers=True):
+def import_set(dset, in_stream, headers=True, **kwargs):
     """Returns dataset from CSV stream."""
 
     dset.wipe()
 
-    if is_py3:
-        rows = csv.reader(StringIO(in_stream))
-    else:
-        rows = csv.reader(StringIO(in_stream), encoding=DEFAULT_ENCODING)
+    if not is_py3:
+        kwargs.setdefault('encoding', DEFAULT_ENCODING)
+
+    rows = csv.reader(StringIO(in_stream), **kwargs)
     for i, row in enumerate(rows):
 
         if (i == 0) and (headers):
diff --git a/tablib/formats/_xlsx.py b/tablib/formats/_xlsx.py
index d697d9c..0cd8500 100644
--- a/tablib/formats/_xlsx.py
+++ b/tablib/formats/_xlsx.py
@@ -33,21 +33,21 @@ def detect(stream):
     except openpyxl.shared.exc.InvalidFileException:
         pass
 
-def export_set(dataset):
+def export_set(dataset, freeze_panes=True):
     """Returns XLSX representation of Dataset."""
 
     wb = Workbook()
     ws = wb.worksheets[0]
     ws.title = dataset.title if dataset.title else 'Tablib Dataset'
 
-    dset_sheet(dataset, ws)
+    dset_sheet(dataset, ws, freeze_panes=freeze_panes)
 
     stream = BytesIO()
     wb.save(stream)
     return stream.getvalue()
 
 
-def export_book(databook):
+def export_book(databook, freeze_panes=True):
     """Returns XLSX representation of DataBook."""
 
     wb = Workbook()
@@ -56,7 +56,7 @@ def export_book(databook):
         ws = wb.create_sheet()
         ws.title = dset.title if dset.title else 'Sheet%s' % (i)
 
-        dset_sheet(dset, ws)
+        dset_sheet(dset, ws, freeze_panes=freeze_panes)
 
 
     stream = BytesIO()
@@ -103,7 +103,7 @@ def import_book(dbook, in_stream, headers=True):
         dbook.add_sheet(data)
 
 
-def dset_sheet(dataset, ws):
+def dset_sheet(dataset, ws, freeze_panes=True):
     """Completes given worksheet from given Dataset."""
     _package = dataset._package(dicts=False)
 
@@ -125,8 +125,9 @@ def dset_sheet(dataset, ws):
                 ws.cell('%s%s'%(col_idx, row_number)).value = unicode(col)
                 style = ws.get_style('%s%s' % (col_idx, row_number))
                 style.font.bold = True
-                ws.freeze_panes = '%s%s' % (frzn_col_idx, row_number)
-
+                if freeze_panes:
+                    # We want to freeze the column after the last column
+                    ws.freeze_panes = '%s%s' % (frzn_col_idx, row_number)
 
             # bold separators
             elif len(row) < dataset.width:
diff --git a/test_tablib.py b/test_tablib.py
index 78d1045..35db877 100755
--- a/test_tablib.py
+++ b/test_tablib.py
@@ -719,5 +719,30 @@ Old       |Man       |100500
         except tablib.InvalidDatasetType:
             self.fail("Subclass of tablib.Dataset should be accepted by Databook.add_sheet")
 
+
+    def test_csv_formatter_support_kwargs(self):
+        """Test CSV import and export with formatter configuration."""
+        data.append(self.john)
+        data.append(self.george)
+        data.headers = self.headers
+
+        expected = 'first_name;last_name;gpa\nJohn;Adams;90\nGeorge;Washington;67\n'
+
+        kwargs = dict(delimiter=';', lineterminator='\n')
+        _csv = data.export('csv', **kwargs)
+        self.assertEqual(expected, _csv)
+
+        # the import works but consider default delimiter=','
+        d1 = tablib.import_set(_csv, format="csv")
+        self.assertEqual(1, len(d1.headers))
+
+        d2 = tablib.import_set(_csv, format="csv", **kwargs)
+        self.assertEqual(3, len(d2.headers))
+
+    def test_databook_formatter_support_kwargs(self):
+        """Test XLSX export with formatter configuration."""
+        self.founders.export('xlsx', freeze_panes=False)
+
+
 if __name__ == '__main__':
     unittest.main()