diff options
Diffstat (limited to 'tablib')
| -rw-r--r-- | tablib/core.py | 106 | ||||
| -rw-r--r-- | tablib/formats/_csv.py | 22 | ||||
| -rw-r--r-- | tablib/formats/_xlsx.py | 15 |
3 files changed, 105 insertions, 38 deletions
diff --git a/tablib/core.py b/tablib/core.py index e15206e..7d09789 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 @@ -248,13 +250,15 @@ 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)) - setattr(cls, 'set_%s' % fmt.title, fmt.import_set) + cls._formats[fmt.title] = (fmt.export_set, None) setattr(cls, 'get_%s' % fmt.title, fmt.export_set) + setattr(cls, 'set_%s' % fmt.title, fmt.import_set) except AttributeError: - pass + cls._formats[fmt.title] = (None, None) def _validate(self, row=None, col=None, safety=False): @@ -424,12 +428,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. @@ -940,7 +966,7 @@ class Dataset(object): return _dset - def unique(self): + def remove_duplicates(self): """Removes all duplicate rows from the :class:`Dataset` object while maintaining the original order.""" seen = set() @@ -961,19 +987,19 @@ class Dataset(object): # Don't return if no data if not self: return - + if rows is None: rows = list(range(self.height)) - + if cols is None: cols = list(self.headers) - + #filter out impossible rows and columns rows = [row for row in rows if row in range(self.height)] cols = [header for header in cols if header in self.headers] _dset = Dataset() - + #filtering rows and columns _dset.headers = list(cols) @@ -989,7 +1015,7 @@ class Dataset(object): if row_no in rows: _dset.append(row=Row(data_row)) - + return _dset @@ -998,6 +1024,8 @@ class Databook(object): """A book of :class:`Dataset` objects. """ + _formats = {} + def __init__(self, sets=None): if sets is None: @@ -1013,7 +1041,6 @@ class Databook(object): except AttributeError: return '<databook object>' - def wipe(self): """Removes all :class:`Dataset` objects from the :class:`Databook`.""" self._datasets = [] @@ -1026,11 +1053,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 @@ -1065,6 +1094,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.""" @@ -1077,32 +1130,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 2be1895..994b23b 100644 --- a/tablib/formats/_csv.py +++ b/tablib/formats/_csv.py @@ -14,14 +14,15 @@ DEFAULT_ENCODING = 'utf-8' DEFAULT_DELIMITER = ',' -def export_set(dataset, delimiter=DEFAULT_DELIMITER): +def export_set(dataset, **kwargs): """Returns CSV representation of Dataset.""" stream = StringIO() - if is_py3: - _csv = csv.writer(stream, delimiter=delimiter) - else: - _csv = csv.writer(stream, delimiter=delimiter, encoding=DEFAULT_ENCODING) + kwargs.setdefault('delimiter', DEFAULT_DELIMITER) + 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 +30,16 @@ def export_set(dataset, delimiter=DEFAULT_DELIMITER): return stream.getvalue() -def import_set(dset, in_stream, headers=True, delimiter=DEFAULT_DELIMITER): +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), delimiter=delimiter) - else: - rows = csv.reader(StringIO(in_stream), delimiter=delimiter, encoding=DEFAULT_ENCODING) + kwargs.setdefault('delimiter', DEFAULT_DELIMITER) + 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 34e83e6..d9d3d57 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) @@ -123,8 +123,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 = 'A2' - + 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: |