def test_index_len(self):
"""Test streamlit.data_frame_proto._index_len."""
# Plain
plain_idx = Index()
plain_idx.plain_index.data.int64s.data.extend([1, 2, 3])
self.assertEqual(3, data_frame_proto._index_len(plain_idx))
# Range
range_idx = Index()
range_idx.range_index.start = 2
range_idx.range_index.stop = 10
self.assertEqual(8, data_frame_proto._index_len(range_idx))
# Multi with no labels
multi_idx = Index()
multi_idx.multi_index.levels.extend([plain_idx, range_idx])
self.assertEqual(0, data_frame_proto._index_len(multi_idx))
# Multi with labels
int32_array = Int32Array()
int32_array.data.extend([4, 5])
multi_idx.multi_index.labels.extend([int32_array])
self.assertEqual(2, data_frame_proto._index_len(multi_idx))
# Datetime
dt_idx = Index()
dt_idx.datetime_index.data.data.extend([1, 2, 3])
self.assertEqual(3, data_frame_proto._index_len(dt_idx))
# TimeDelta
td_idx = Index()
td_idx.timedelta_index.data.data.extend([1, 2, 3, 4])
self.assertEqual(4, data_frame_proto._index_len(td_idx))
# Ine64
i64_idx = Index()
i64_idx.int_64_index.data.data.extend([1, 2, 3, 4, 5])
self.assertEqual(5, data_frame_proto._index_len(i64_idx))
# Float64
f64_idx = Index()
f64_idx.float_64_index.data.data.extend([1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
self.assertEqual(6, data_frame_proto._index_len(f64_idx))
def test_marshall_index(self):
"""Test streamlit.data_frame_proto._marshall_index."""
df = pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]})
# Plain Index
proto = Index()
data_frame_proto._marshall_index(df.columns, proto)
self.assertEqual(["col1", "col2"], proto.plain_index.data.strings.data)
# Range Index
proto = Index()
data_frame_proto._marshall_index(df.index, proto)
self.assertEqual(0, proto.range_index.start)
self.assertEqual(2, proto.range_index.stop)
# Range Index with NaNs
df_nan = pd.DataFrame(data={"col1": [], "col2": []})
proto = Index()
data_frame_proto._marshall_index(df_nan.index, proto)
self.assertEqual(0, proto.range_index.start)
self.assertEqual(0, proto.range_index.stop)
# multi index
df_multi = pd.MultiIndex.from_arrays([[1, 2], [3, 4]], names=["one", "two"])
proto = Index()
data_frame_proto._marshall_index(df_multi, proto)
self.assertEqual([1, 2], proto.multi_index.levels[0].int_64_index.data.data)
self.assertEqual([0, 1], proto.multi_index.labels[0].data)
# datetimeindex
truth = [int(x * 1e9) for x in (1554138000, 1554141600, 1554145200)]
df_dt = pd.date_range(
start="2019/04/01 10:00", end="2019/04/01 12:00", freq="H"
)
proto = Index()
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_index(df_dt, proto)
self.assertEqual(truth, proto.datetime_index.data.data)
# timedeltaindex
df_td = pd.to_timedelta(np.arange(1, 5), unit="ns")
proto = Index()
data_frame_proto._marshall_index(df_td, proto)
self.assertEqual([1, 2, 3, 4], proto.timedelta_index.data.data)
# int64index
df_int64 = pd.Int64Index(np.arange(1, 5))
proto = Index()
data_frame_proto._marshall_index(df_int64, proto)
self.assertEqual([1, 2, 3, 4], proto.int_64_index.data.data)
# float64index
df_float64 = pd.Float64Index(np.arange(1, 5))
proto = Index()
data_frame_proto._marshall_index(df_float64, proto)
self.assertEqual([1, 2, 3, 4], proto.float_64_index.data.data)
# Period index
df_period = pd.period_range(
start="2005-12-21 08:45 ", end="2005-12-21 11:55", freq="H"
)
proto = Index()
with pytest.raises(NotImplementedError) as e:
data_frame_proto._marshall_index(df_period, proto)
err_msg = (
"Can't handle <class 'pandas.core.indexes.period.PeriodIndex'>" " yet."
)
self.assertEqual(err_msg, str(e.value))
def test_concat_index(self):
"""Test streamlit.data_frame_proto._concat_index."""
# Empty
idx0 = Index()
idx0.plain_index.data.int64s.data.extend([])
idx1 = Index()
idx1.plain_index.data.int64s.data.extend([1, 2])
data_frame_proto._concat_index(idx0, idx1)
self.assertEqual(idx0, idx1)
# type mismatch
idx2 = Index()
idx2.plain_index.data.doubles.data.extend([3.0, 4.0])
with pytest.raises(ValueError) as e:
data_frame_proto._concat_index(idx1, idx2)
err_msg = "Cannot concatenate int64s with doubles."
self.assertEqual(err_msg, str(e.value))
# plain index
idx3 = Index()
idx3.plain_index.data.int64s.data.extend([5, 6])
idx4 = Index()
idx4.plain_index.data.int64s.data.extend([1, 2, 5, 6])
data_frame_proto._concat_index(idx1, idx3)
self.assertEqual(idx1, idx4)
# range index
r_idx1 = Index()
r_idx1.range_index.start = 2
r_idx1.range_index.stop = 10
r_idx2 = Index()
r_idx2.range_index.start = 10
r_idx2.range_index.stop = 20
r_combined = Index()
r_combined.range_index.start = 2
r_combined.range_index.stop = 20
data_frame_proto._concat_index(r_idx1, r_idx2)
self.assertEqual(r_idx1, r_combined)
# multi
int32_array = Int32Array()
int32_array.data.extend([4, 5])
m_idx1 = Index()
m_idx1.multi_index.labels.extend([int32_array])
m_idx2 = Index()
m_idx2.multi_index.labels.extend([int32_array])
with pytest.raises(NotImplementedError) as e:
data_frame_proto._concat_index(m_idx1, m_idx2)
err_msg = "Cannot yet concatenate MultiIndices."
self.assertEqual(err_msg, str(e.value))
# int_64_index
i_idx1 = Index()
i_idx1.int_64_index.data.data.extend([1, 2])
i_idx2 = Index()
i_idx2.int_64_index.data.data.extend([3, 4])
i_combined = Index()
i_combined.int_64_index.data.data.extend([1, 2, 3, 4])
data_frame_proto._concat_index(i_idx1, i_idx2)
self.assertEqual(i_idx1, i_combined)
# datetime_index
dt_idx1 = Index()
dt_idx1.datetime_index.data.data.extend([1, 2])
dt_idx2 = Index()
dt_idx2.datetime_index.data.data.extend([3, 4])
dt_combined = Index()
dt_combined.datetime_index.data.data.extend([1, 2, 3, 4])
data_frame_proto._concat_index(dt_idx1, dt_idx2)
self.assertEqual(dt_idx1, dt_combined)
# timedelta_index
td_idx1 = Index()
td_idx1.timedelta_index.data.data.extend([1, 2])
td_idx2 = Index()
td_idx2.timedelta_index.data.data.extend([3, 4])
td_combined = Index()
td_combined.timedelta_index.data.data.extend([1, 2, 3, 4])
data_frame_proto._concat_index(td_idx1, td_idx2)
self.assertEqual(td_idx1, td_combined)
# Not implemented
f_idx1 = Index()
f_idx1.float_64_index.data.data.extend([1.0, 2.0])
f_idx2 = Index()
f_idx2.float_64_index.data.data.extend([3.0, 4.0])
with pytest.raises(NotImplementedError) as e:
data_frame_proto._concat_index(f_idx1, f_idx2)
err_msg = 'Cannot concatenate "float_64_index" indices.'
self.assertEqual(err_msg, str(e.value))