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))