def export(path, rows, exclude=None, include=None, update=None, append=False, header=True):
path = pathlib.Path(path)
if path.suffix == '.txt':
Writer = TxtWriter
elif path.suffix == '.tsv':
Writer = functools.partial(CsvWriter, dialect='excel-tab', lineterminator='\n')
elif path.suffix == '.csv':
Writer = functools.partial(CsvWriter, dialect='excel', lineterminator='\n')
else:
raise ValueError("Unknown file format '%s'." % path.suffix)
mode = 'a' if append else 'w'
with path.open(mode) as f:
writer = Writer(f)
rows = flatten(rows, exclude, include, update)
for row in rows:
if header:
writer.writeheader(row)
break
for row in rows:
writer.writerow(row)
def test_flatten_rows_include(data):
rows = [
Row(key=1, value={'a': 1}),
Row(key=2, value={'b': 2}),
]
assert list(flatten(rows, include=['a', 'b'])) == [
('a', 'b'),
(1, None),
(None, 2),
]
def test_flatten_int_key(data):
rows = [
Row(key=1, value={'year': {2000: 1, 2001: 2}}),
Row(key=2, value={'year': {2000: 3, 2001: 4}}),
]
assert list(flatten(rows)) == [
('key', 'year.2000', 'year.2001'),
(1, 1, 2),
(2, 3, 4),
]
def test_flatten_rows_value(data):
rows = [
Row(key=1, value='a'),
Row(key=2, value='b'),
]
assert list(flatten(rows)) == [
('key', 'value'),
(1, 'a'),
(2, 'b'),
]
def test_flatten_rows_include(data):
rows = [
Row(key=1, value={'a': 1}),
Row(key=2, value={'b': 2}),
]
assert list(flatten(rows, include=['a', 'b'])) == [
('a', 'b'),
(1, None),
(None, 2),
]
def test_flatten_rows_value(data):
rows = [
Row(key=1, value='a'),
Row(key=2, value='b'),
]
assert list(flatten(rows)) == [
('key', 'value'),
(1, 'a'),
(2, 'b'),
]
def test_flatten_rows_update_without_include(data):
rows = [
Row(key=1, value={'text': 'abc'}),
Row(key=1, value={'text': 'abcde'}),
]
update = {'size': databot.this.value.text.apply(len)}
assert list(flatten(rows, update=update)) == [
('key', 'size', 'text'),
(1, 3, 'abc'),
(1, 5, 'abcde'),
]
def test_flatten_rows_update(data):
rows = [
Row(key=1, value={'text': 'abc'}),
Row(key=1, value={'text': 'abcde'}),
]
update = {'size': databot.this.value.text.apply(len)}
assert list(flatten(rows, include=['key', 'size'], update=update)) == [
('key', 'size'),
(1, 3),
(1, 5),
]
def test_flatten_rows_callable_update(data):
rows = [
Row(key=1, value={'text': 'abc'}),
Row(key=1, value={'text': 'abcde'}),
]
def update(row):
return {'size': len(row.value['text'])}
assert list(flatten(rows, update=update)) == [
('size', ),
(3, ),
(5, ),
]
def test_flatten_rows_callable_update(data):
rows = [
Row(key=1, value={'text': 'abc'}),
Row(key=1, value={'text': 'abcde'}),
]
def update(row):
return {'size': len(row.value['text'])}
assert list(flatten(rows, update=update)) == [
('size',),
(3,),
(5,),
]
def test_flatten():
rows = [
Row(key=1, value={'foo': 'bar', 'events': [
{'name': 'Event 1', 'date': '2017-01-01', 'people': ['a', 'b']},
{'name': 'Event 2', 'date': '2017-01-02', 'people': ['a']},
]}),
Row(key=2, value={'foo': 'baz', 'events': [
{'name': 'Event 3', 'date': '2017-01-03', 'people': ['x', 'y']},
{'name': 'Event 4', 'date': '2017-01-04', 'people': ['z']},
]}),
]
assert list(flatten(rows)) == [
('events.date', 'events.name', 'events.people', 'foo', 'key'),
('2017-01-01', 'Event 1', 'a', 'bar', 1),
('2017-01-01', 'Event 1', 'b', 'bar', 1),
('2017-01-02', 'Event 2', 'a', 'bar', 1),
('2017-01-03', 'Event 3', 'x', 'baz', 2),
('2017-01-03', 'Event 3', 'y', 'baz', 2),
('2017-01-04', 'Event 4', 'z', 'baz', 2),
]
assert list(flatten(rows, include=('key', 'foo', 'events.people'))) == [
('key', 'foo', 'events.people'),
(1, 'bar', 'a'),
(1, 'bar', 'b'),
(1, 'bar', 'a'),
(2, 'baz', 'x'),
(2, 'baz', 'y'),
(2, 'baz', 'z'),
]
assert list(flatten(rows, include=('key', 'foo'))) == [
('key', 'foo'),
(1, 'bar'),
(2, 'baz'),
]
def test_flatten_int_key(data):
rows = [
Row(key=1, value={'year': {
2000: 1,
2001: 2
}}),
Row(key=2, value={'year': {
2000: 3,
2001: 4
}}),
]
assert list(flatten(rows)) == [
('key', 'year.2000', 'year.2001'),
(1, 1, 2),
(2, 3, 4),
]
def test_flatten_list(data):
rows = [
Row(key=1, value={'events': [
{'name': 'Event 1', 'date': '2017-01-01'},
{'name': 'Event 2', 'date': '2017-02-01'},
]}),
Row(key=2, value={'events': [
{'name': 'Event 3', 'date': '2017-03-01'},
{'name': 'Event 4', 'date': '2017-04-01'},
]}),
]
assert list(flatten(rows)) == [
('events.date', 'events.name', 'key'),
('2017-01-01', 'Event 1', 1),
('2017-02-01', 'Event 2', 1),
('2017-03-01', 'Event 3', 2),
('2017-04-01', 'Event 4', 2),
]
def table(self, rows, exclude=None, include=None):
_rows = []
flat_rows = flatten(rows, exclude, include)
for row in flat_rows:
max_value_size = (self.width // len(row)) * 3
_rows.append(row)
break
for row in flat_rows:
_row = []
for value in row:
if isinstance(value, list):
value = repr(value)
if isinstance(value, str) and len(value) > max_value_size:
value = value[:max_value_size] + '...'
_row.append(value)
_rows.append(_row)
table = texttable.Texttable(self.width)
table.set_deco(texttable.Texttable.HEADER)
table.add_rows(_rows)
self.info(table.draw())
def test_flatten_list(data):
rows = [
Row(key=1,
value={
'events': [
{
'name': 'Event 1',
'date': '2017-01-01'
},
{
'name': 'Event 2',
'date': '2017-02-01'
},
]
}),
Row(key=2,
value={
'events': [
{
'name': 'Event 3',
'date': '2017-03-01'
},
{
'name': 'Event 4',
'date': '2017-04-01'
},
]
}),
]
assert list(flatten(rows)) == [
('events.date', 'events.name', 'key'),
('2017-01-01', 'Event 1', 1),
('2017-02-01', 'Event 2', 1),
('2017-03-01', 'Event 3', 2),
('2017-04-01', 'Event 4', 2),
]
def export(pd, rows, exclude=None, include=None, update=None):
rows = flatten(rows, exclude, include, update)
columns = next(rows, None)
index, columns = index_from_columns(columns, 'key')
items = rows_to_dataframe_items(rows, index)
return pd.DataFrame.from_items(items, columns=columns, orient='index')
def test_flatten():
rows = [
Row(key=1,
value={
'foo':
'bar',
'events': [
{
'name': 'Event 1',
'date': '2017-01-01',
'people': ['a', 'b']
},
{
'name': 'Event 2',
'date': '2017-01-02',
'people': ['a']
},
]
}),
Row(key=2,
value={
'foo':
'baz',
'events': [
{
'name': 'Event 3',
'date': '2017-01-03',
'people': ['x', 'y']
},
{
'name': 'Event 4',
'date': '2017-01-04',
'people': ['z']
},
]
}),
]
assert list(flatten(rows)) == [
('events.date', 'events.name', 'events.people', 'foo', 'key'),
('2017-01-01', 'Event 1', 'a', 'bar', 1),
('2017-01-01', 'Event 1', 'b', 'bar', 1),
('2017-01-02', 'Event 2', 'a', 'bar', 1),
('2017-01-03', 'Event 3', 'x', 'baz', 2),
('2017-01-03', 'Event 3', 'y', 'baz', 2),
('2017-01-04', 'Event 4', 'z', 'baz', 2),
]
assert list(flatten(rows, include=('key', 'foo', 'events.people'))) == [
('key', 'foo', 'events.people'),
(1, 'bar', 'a'),
(1, 'bar', 'b'),
(1, 'bar', 'a'),
(2, 'baz', 'x'),
(2, 'baz', 'y'),
(2, 'baz', 'z'),
]
assert list(flatten(rows, include=('key', 'foo'))) == [
('key', 'foo'),
(1, 'bar'),
(2, 'baz'),
]