mirror of
https://github.com/kennethreitz/tablib.git
synced 2026-06-05 23:10:17 +00:00
New import/export interface with dataset and databook import_ and export methods
and overloaded `import_set` and `import_book` functions.
This commit is contained in:
+78
-14
@@ -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"
|
||||
|
||||
|
||||
+10
-10
@@ -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):
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user