def test_begin_twice():
transaction = Transaction(Collection([]))
transaction.begin()
raises(TransactionAlreadyActiveError, transaction.begin)
# no exception raised if transaction is already finished
transaction.rollback()
transaction.begin()
class Collection(object):
"""Basic data-collection.
Example:
>>> Collection(((1,2,3), (2,3,4)))
<Collection 2 rows, 3 columns>
>>> Collection([])
<Collection Empty>
"""
def __init__(self, data, **kwargs):
self.data = [list(x) for x in data]
self.width = 0 if not self.data else len(self.data[0])
self.transaction = Transaction(self)
# State vars
self._child_collections = {}
# Add filters
def _common_kwarg_handling():
if 'filter' in kwargs:
for filter_fn in kwargs['filter']:
self.filter(filter_fn)
if 'group' in kwargs:
self.group(kwargs['group'])
self._handle_kwargs(_common_kwarg_handling, **kwargs)
def __len__(self):
return len(self.data)
def __iter__(self):
for idx, record in enumerate(self.data):
if idx in self._child_collections:
children = self._child_collections[idx]
else:
children = None
yield Record(record, children)
def __getitem__(self, key):
if key in self._child_collections:
return Record(self.data[key], self._child_collections[key])
else:
return Record(self.data[key])
def __enter__(self):
self.transaction.begin()
def __exit__(self, exc_type, exc_value, traceback):
if not exc_type:
self.transaction.commit()
else:
self.transaction.rollback()
def __repr__(self):
if self.data:
return ("<Collection %s rows, %s columns>"
% (len(self.data), len(self.data[0])))
else:
return "<Collection Empty>"
def add_formatted_column(self, fmt):
"""Add new column defined by given format string.
The python format specification is used to create the resulting
columns. (http://docs.python.org/library/string.html#formatstrings)
>>> col = Collection((('a', 'b'), ('c', 'd')))
>>> col.add_formatted_column('{0}, {1}')
>>> col[0][2]
'a, b'
>>> col[1][2]
'c, d'
"""
def _do_format(row, collection):
return fmt.format(*row)
self.transaction.add('new_cols', _do_format)
def add_calculated_column(self, calculation):
"""Add new column whose value is the result of the given calculation."""
# replace place-holders with dictionary refs
calculation = calculation.replace('{', 'r[').replace('}', ']')
# create new function and return
exec '_do_calc = lambda r, c: ' + calculation
self.transaction.add('new_cols', _do_calc)
def factory(self, data):
"""Returns method to generate similar collection instance."""
return type(self)(data)
def filter(self, fn):
"""Filter rows that match given function."""
self.transaction.add('filter', fn)
def group(self, groupby_list):
"""Group rows with the same values for the given column positions."""
for groupby in groupby_list:
self.transaction.add('group', groupby)
def _handle_kwargs(self, common_kwarg_handling, **kwargs):
"""Handle kwargs passed in on __init__."""
with self:
common_kwarg_handling()
if 'coerce' in kwargs:
for i in xrange(len(self.data)):
for idx, type_ in kwargs['coerce'].iteritems():
self.data[i][idx] = type_(self.data[i][idx])
if 'formatted_columns' in kwargs:
for col in kwargs['formatted_columns']:
self.add_formatted_column(col)
if 'calculated_columns' in kwargs:
for col in kwargs['calculated_columns']:
self.add_calculated_column(col)