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 = 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 = Delta()
data_frame_proto._get_data_frame(delta)
err_msg = "Cannot extract DataFrame from None."
self.assertEqual(err_msg, str(e.value))
# Generic Data
aa = AnyArray()
aa.int64s.data.extend([1, 2, 3])
# Delta DataFrame
delta_df = 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 = 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)
# Vega-Lite Chart
delta_vega = 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 = 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 = 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 = 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 = 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_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 = 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 = 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 = AnyArray()
aa.int64s.data.extend([1, 2, 3])
# Delta DataFrame
delta_df = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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))
