mirror of
https://github.com/kennethreitz/tablib.git
synced 2026-06-05 23:10:17 +00:00
adding support for creating subset of a dataset.
This commit is contained in:
@@ -28,3 +28,4 @@ Patches and Suggestions
|
||||
- Marc Abramowitz
|
||||
- Alex Gaynor
|
||||
- James Douglass
|
||||
- Rabin Nankhwa
|
||||
@@ -933,6 +933,42 @@ class Dataset(object):
|
||||
self.__headers = None
|
||||
|
||||
|
||||
def subset(self, rows=None, cols=None):
|
||||
"""Returns a new instance of the :class:`Dataset`,
|
||||
including only specified rows and columns.
|
||||
"""
|
||||
|
||||
# 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)
|
||||
|
||||
_dset = Dataset()
|
||||
|
||||
#filtering rows and columns
|
||||
_dset.headers = list(cols)
|
||||
|
||||
_dset._data = []
|
||||
for row_no, row in enumerate(self._data):
|
||||
data_row = []
|
||||
for key in _dset.headers:
|
||||
if key in self.headers:
|
||||
pos = self.headers.index(key)
|
||||
data_row.append(row[pos])
|
||||
else:
|
||||
raise KeyError
|
||||
|
||||
if row_no in rows:
|
||||
_dset.append(row=Row(data_row))
|
||||
|
||||
return _dset
|
||||
|
||||
|
||||
|
||||
class Databook(object):
|
||||
"""A book of :class:`Dataset` objects.
|
||||
|
||||
@@ -713,6 +713,26 @@ class TablibTestCase(unittest.TestCase):
|
||||
self.assertTrue(data[0] == new_row)
|
||||
|
||||
|
||||
def test_subset(self):
|
||||
"""Create a subset of a dataset"""
|
||||
|
||||
rows = (0, 2)
|
||||
columns = ('first_name','gpa')
|
||||
|
||||
data.headers = self.headers
|
||||
|
||||
data.append(self.john)
|
||||
data.append(self.george)
|
||||
data.append(self.tom)
|
||||
|
||||
#Verify data is truncated
|
||||
subset = data.subset(rows=rows, cols=columns)
|
||||
self.assertEqual(type(subset), tablib.Dataset)
|
||||
self.assertEqual(subset.headers, list(columns))
|
||||
self.assertEqual(subset._data[0].list, ['John', 90])
|
||||
self.assertEqual(subset._data[1].list, ['Thomas', 50])
|
||||
|
||||
|
||||
def test_formatters(self):
|
||||
"""Confirm formatters are being triggered."""
|
||||
|
||||
|
||||
Reference in New Issue
Block a user