def test_concat_any_array(self):
"""Test streamlit.data_frame_proto._concat_any_array."""
aa0 = protobuf.AnyArray()
aa0.int64s.data.extend([])
aa1 = protobuf.AnyArray()
aa1.int64s.data.extend([1, 2])
aa2 = protobuf.AnyArray()
aa2.int64s.data.extend([3, 4])
aa3 = protobuf.AnyArray()
aa3.doubles.data.extend([5.0, 6.0])
combined = protobuf.AnyArray()
combined.int64s.data.extend([1, 2, 3, 4])
# both not empty
data_frame_proto._concat_any_array(aa1, aa2)
self.assertEqual(aa1, combined)
# empty
data_frame_proto._concat_any_array(aa0, aa1)
self.assertEqual(aa0, aa1)
# types dont match
with pytest.raises(ValueError) as e:
data_frame_proto._concat_any_array(aa2, aa3)
err_msg = 'Cannot concatenate int64s with doubles.'
self.assertEqual(err_msg, str(e.value))
def test_any_array_len(self):
"""Test streamlit.data_frame_proto._any_array_len."""
data = [
('strings', 2, ['a', 'b']),
('int64s', 3, [1, 2, 3]),
('doubles', 4, [1.0, 2.0, 3.0, 4.0]),
# datetimes and timedeltas are just stored as ints and aren't
# python data types.
('datetimes', 5, [1, 2, 3, 4, 5]),
('timedeltas', 6, [1, 2, 3, 4, 5, 6]),
]
for kind, length, array in data:
aa = protobuf.AnyArray()
pb = getattr(aa, kind)
pb.data.extend(array)
self.assertEqual(length, data_frame_proto._any_array_len(aa))
def test_get_or_create_dataset(self):
"""Test streamlit.data_frame_proto._get_or_create_dataset."""
chart = VegaLiteChart()
ds1 = NamedDataSet()
ds1.name = 'dataset 1'
ds1.has_name = True
aa = protobuf.AnyArray()
aa.int64s.data.extend([1, 2, 3])
ds1.data.data.cols.extend([aa])
ds2 = NamedDataSet()
ds2.name = 'dataset 2'
ds2.has_name = True
chart.datasets.extend([ds1, ds2])
ret = data_frame_proto._get_or_create_dataset(chart.datasets,
'dataset 1')
self.assertEqual(ret, ds1.data)
def test_get_data_frame(self):
"""Test streamlit.data_frame_proto._get_data_frame."""
# Test delta not new_element or add_rows
with pytest.raises(ValueError) as e:
delta = protobuf.Delta()
data_frame_proto._get_data_frame(delta)
err_msg = 'Cannot extract DataFrame from None.'
self.assertEqual(err_msg, str(e.value))
# Test delta = new_element, a name is used and type is chart, df, table
with pytest.raises(ValueError) as e:
delta = protobuf.Delta()
# TODO(armando): test df and table
delta.new_element.chart.type = 'some chart'
data_frame_proto._get_data_frame(delta, name='some name')
err_msg = 'Dataset names not supported for st.chart'
self.assertEqual(err_msg, str(e.value))
# Generic Data
aa = protobuf.AnyArray()
aa.int64s.data.extend([1, 2, 3])
# Delta DataFrame
delta_df = protobuf.Delta()
delta_df.new_element.data_frame.data.cols.extend([aa])
df = data_frame_proto._get_data_frame(delta_df)
self.assertEqual(df, delta_df.new_element.data_frame)
# Delta Table
delta_table = protobuf.Delta()
delta_table.new_element.table.data.cols.extend([aa])
df = data_frame_proto._get_data_frame(delta_table)
self.assertEqual(df, delta_table.new_element.table)
# Delta Chart
delta_chart = protobuf.Delta()
delta_chart.new_element.chart.data.data.cols.extend([aa])
df = data_frame_proto._get_data_frame(delta_chart)
self.assertEqual(df, delta_chart.new_element.chart.data)
# Vega-Lite Chart
delta_vega = protobuf.Delta()
delta_vega.new_element.vega_lite_chart.data.data.cols.extend([aa])
df = data_frame_proto._get_data_frame(delta_vega)
self.assertEqual(df, delta_vega.new_element.vega_lite_chart.data)
# Vega-Lite Chart w/ named dataset
delta_vega_dataset = protobuf.Delta()
ds1 = NamedDataSet()
ds1.name = 'dataset 1'
ds1.has_name = True
ds1.data.data.cols.extend([aa])
delta_vega_dataset.new_element.vega_lite_chart.datasets.extend([ds1])
df = data_frame_proto._get_data_frame(delta_vega_dataset, 'dataset 1')
self.assertEqual(
df,
delta_vega_dataset.new_element.vega_lite_chart.datasets[0].data)
# Vega-Lite Chart w/ unnamed dataset
delta_vega_unnamed_dataset = protobuf.Delta()
ds2 = NamedDataSet()
ds2.has_name = False
ds2.data.data.cols.extend([aa])
delta_vega_unnamed_dataset.new_element.vega_lite_chart.datasets.extend(
[ds2])
df = data_frame_proto._get_data_frame(delta_vega_unnamed_dataset)
self.assertEqual(
df, delta_vega_unnamed_dataset.new_element.vega_lite_chart.
datasets[0].data)
# add_rows w/ name
delta_add_rows = protobuf.Delta()
delta_add_rows.add_rows.name = 'named dataset'
delta_add_rows.add_rows.has_name = True
delta_add_rows.add_rows.data.data.cols.extend([aa])
df = data_frame_proto._get_data_frame(delta_add_rows, 'named dataset')
self.assertEqual(df, delta_add_rows.add_rows.data)
# add_rows w/out name
with pytest.raises(ValueError) as e:
delta_add_rows_noname = protobuf.Delta()
delta_add_rows_noname.add_rows.name = 'named dataset'
delta_add_rows_noname.add_rows.has_name = True
delta_add_rows_noname.add_rows.data.data.cols.extend([aa])
df = data_frame_proto._get_data_frame(delta_add_rows_noname)
err_msg = 'No dataset found with name "None".'
self.assertEqual(err_msg, str(e.value))
def test_add_rows(self):
"""Test streamlit.data_frame_proto._add_rows."""
# Generic Data
aa = protobuf.AnyArray()
aa.int64s.data.extend([1, 2])
cell_style = CellStyle()
cell_style.css.extend([_css_style('color', 'black')])
style = CellStyleArray()
style.styles.extend([cell_style])
# Delta DataFrame
dt1 = protobuf.Delta()
dt1.new_element.data_frame.data.cols.extend([aa])
dt1.new_element.data_frame.index.plain_index.data.int64s.data.extend(
[3, 4])
dt1.new_element.data_frame.columns.plain_index.data.int64s.data.extend(
[5, 6])
dt1.new_element.data_frame.style.cols.extend([style])
dt2 = protobuf.Delta()
dt2.new_element.data_frame.data.cols.extend([aa])
dt2.new_element.data_frame.index.plain_index.data.int64s.data.extend(
[3, 4])
dt2.new_element.data_frame.columns.plain_index.data.int64s.data.extend(
[5, 6])
dt2.new_element.data_frame.style.cols.extend([style])
combined = protobuf.Delta()
aa_combined = protobuf.AnyArray()
aa_combined.int64s.data.extend([1, 2, 1, 2])
style_combined = CellStyleArray()
style_combined.styles.extend([cell_style, cell_style])
combined.new_element.data_frame.data.cols.extend([aa_combined])
row_index = combined.new_element.data_frame.index.plain_index
row_index.data.int64s.data.extend([3, 4, 3, 4])
col_index = combined.new_element.data_frame.columns.plain_index
col_index.data.int64s.data.extend([5, 6])
combined.new_element.data_frame.style.cols.extend([style_combined])
# Test both not empty
data_frame_proto.add_rows(dt1, dt2)
self.assertEqual(dt1, combined)
# Test one empty
dt0 = protobuf.Delta()
dt0.new_element.data_frame.data.cols.extend([])
data_frame_proto.add_rows(dt0, dt1)
self.assertEqual(str(dt0), str(dt1))
# Test both empty
empty0 = protobuf.Delta()
empty0.new_element.data_frame.data.cols.extend([])
empty1 = protobuf.Delta()
empty1.new_element.data_frame.data.cols.extend([])
data_frame_proto.add_rows(empty0, empty1)
self.assertEqual(str(empty0), str(empty1))
# Test different data shapes
diff0 = protobuf.Delta()
diff0.new_element.data_frame.data.cols.extend([aa, aa])
diff1 = protobuf.Delta()
diff1.new_element.data_frame.data.cols.extend([aa])
with pytest.raises(ValueError) as e:
data_frame_proto.add_rows(diff0, diff1)
err_msg = 'Dataframes have incompatible shapes'
self.assertEqual(err_msg, str(e.value))
def test_marshall_any_array(self):
"""Test streamlit.data_frame_proto._marshall_any_array."""
# list
list_data = [1, 2]
list_proto = protobuf.AnyArray()
data_frame_proto._marshall_any_array(list_data, list_proto)
self.assertEqual(list_proto.int64s.data, list_data)
# wrong shape
with pytest.raises(ValueError) as e:
data_frame_proto._marshall_any_array([[1, 2], [3, 4]],
protobuf.AnyArray())
err_msg = 'Array must be 1D.'
self.assertEqual(err_msg, str(e.value))
# float
float_data = pd.Series(np.array([1.0, 2.0]), dtype=np.floating)
float_proto = protobuf.AnyArray()
data_frame_proto._marshall_any_array(float_data, float_proto)
self.assertEqual(float_proto.doubles.data, float_data.tolist())
# timedelta64
td_data = np.array([1, 2], dtype=np.timedelta64)
td_proto = protobuf.AnyArray()
data_frame_proto._marshall_any_array(td_data, td_proto)
self.assertEqual(td_proto.timedeltas.data, td_data.tolist())
# int
int_data = np.array([1, 2], dtype=np.integer)
int_proto = protobuf.AnyArray()
data_frame_proto._marshall_any_array(int_data, int_proto)
self.assertEqual(int_proto.int64s.data, int_data.tolist())
# bool
bool_data = np.array([True, False], dtype=np.bool)
bool_proto = protobuf.AnyArray()
data_frame_proto._marshall_any_array(bool_data, bool_proto)
self.assertEqual(bool_proto.int64s.data, bool_data.tolist())
# object
obj_data = np.array([json.dumps, json.dumps], dtype=np.object)
obj_proto = protobuf.AnyArray()
truth = [str(json.dumps), str(json.dumps)]
data_frame_proto._marshall_any_array(obj_data, obj_proto)
self.assertEqual(obj_proto.strings.data, truth)
# No timezone
dt_data = pd.Series([np.datetime64('2019-04-09T12:34:56')])
dt_proto = protobuf.AnyArray()
obj_to_patch = (
'streamlit.elements.data_frame_proto.tzlocal.get_localzone')
with patch(obj_to_patch) as p:
p.return_value = 'America/Los_Angeles'
data_frame_proto._marshall_any_array(dt_data, dt_proto)
self.assertEqual(1554838496.0,
dt_proto.datetimes.data[0] / 1000000000)
# With timezone
dt_data = pd.Series([np.datetime64('2019-04-09T12:34:56')])
dt_data = dt_data.dt.tz_localize('UTC')
data_frame_proto._marshall_any_array(dt_data, dt_proto)
self.assertEqual(1554838496.0, dt_proto.datetimes.data[0] / 1000000000)
# string
str_data = np.array(['random', 'string'])
str_proto = protobuf.AnyArray()
with pytest.raises(NotImplementedError) as e:
data_frame_proto._marshall_any_array(str_data, str_proto)
if sys.version_info >= (3, 0):
err_msg = 'Dtype <U6 not understood.'
else:
err_msg = 'Dtype |S6 not understood.'
self.assertEqual(err_msg, str(e.value))