/test/components/test_dataframe.py

from datetime import datetime

import numpy as np
import pandas as pd
import polars as pl
import pytest

import gradio as gr
from gradio.components.dataframe import DataframeData


class TestDataframe:
    def test_component_functions(self):
        """
        Preprocess, serialize, get_config
        """
        x_data = {
            "data": [["Tim", 12, False], ["Jan", 24, True]],
            "headers": ["Name", "Age", "Member"],
            "metadata": None,
        }
        x_payload = DataframeData(**x_data)  # type: ignore
        dataframe_input = gr.Dataframe(headers=["Name", "Age", "Member"])
        output = dataframe_input.preprocess(x_payload)
        assert output["Age"][1] == 24  # type: ignore
        assert not output["Member"][0]  # type: ignore
        assert dataframe_input.postprocess(output) == x_payload

        dataframe_input = gr.Dataframe(
            headers=["Name", "Age", "Member"], label="Dataframe Input"
        )
        assert dataframe_input.get_config() == {
            "value": {
                "headers": ["Name", "Age", "Member"],
                "data": [],
                "metadata": None,
            },
            "_selectable": False,
            "key": None,
            "preserved_by_key": ["value"],
            "headers": ["Name", "Age", "Member"],
            "row_count": [3, "dynamic"],
            "row_limits": None,
            "col_count": [3, "dynamic"],
            "column_count": [3, "dynamic"],
            "column_limits": None,
            "datatype": "str",
            "type": "pandas",
            "label": "Dataframe Input",
            "show_label": True,
            "scale": None,
            "min_width": 160,
            "interactive": None,
            "visible": True,
            "elem_id": None,
            "elem_classes": [],
            "show_row_numbers": False,
            "show_search": "none",
            "static_columns": [],
            "pinned_columns": None,
            "wrap": False,
            "proxy_url": None,
            "name": "dataframe",
            "max_height": 500,
            "latex_delimiters": [{"display": True, "left": "$$", "right": "$$"}],
            "line_breaks": True,
            "column_widths": [],
            "buttons": None,
            "max_chars": None,
        }
        dataframe_input = gr.Dataframe()
        output = dataframe_input.preprocess(DataframeData(**x_data))  # type: ignore
        assert output["Age"][1] == 24  # type: ignore

        x_data = {
            "data": [["Tim", 12, False], ["Jan", 24, True]],
            "headers": ["Name", "Age", "Member"],
            "metadata": {"display_value": None, "styling": None},
        }
        dataframe_input.preprocess(DataframeData(**x_data))  # type: ignore

        with pytest.raises(ValueError):
            gr.Dataframe(type="unknown")  # type: ignore

        dataframe_output = gr.Dataframe()
        assert dataframe_output.get_config() == {
            "value": {
                "headers": ["1", "2", "3"],
                "data": [],
                "metadata": None,
            },
            "_selectable": False,
            "key": None,
            "preserved_by_key": ["value"],
            "headers": ["1", "2", "3"],
            "row_count": [3, "dynamic"],
            "row_limits": None,
            "col_count": [3, "dynamic"],
            "column_count": [3, "dynamic"],
            "column_limits": None,
            "datatype": "str",
            "type": "pandas",
            "label": None,
            "show_label": True,
            "show_row_numbers": False,
            "show_search": "none",
            "static_columns": [],
            "pinned_columns": None,
            "scale": None,
            "min_width": 160,
            "interactive": None,
            "visible": True,
            "elem_id": None,
            "elem_classes": [],
            "wrap": False,
            "proxy_url": None,
            "name": "dataframe",
            "max_height": 500,
            "latex_delimiters": [{"display": True, "left": "$$", "right": "$$"}],
            "line_breaks": True,
            "column_widths": [],
            "buttons": None,
            "max_chars": None,
        }

        dataframe_input = gr.Dataframe(column_widths=["100px", 200, "50%"])
        assert dataframe_input.get_config()["column_widths"] == [
            "100px",
            "200px",
            "50%",
        ]

    def test_postprocess(self):
        """
        postprocess
        """
        dataframe_output = gr.Dataframe()
        output = dataframe_output.postprocess(np.zeros((2, 2))).model_dump()
        assert output == {
            "data": [[0, 0], [0, 0]],
            "headers": ["1", "2"],
            "metadata": None,
        }
        output = dataframe_output.postprocess([[1, 3, 5]]).model_dump()
        assert output == {
            "data": [[1, 3, 5]],
            "headers": ["1", "2", "3"],
            "metadata": None,
        }
        output = dataframe_output.postprocess(
            pd.DataFrame([[2, True], [3, True], [4, False]], columns=["num", "prime"])  # type: ignore
        ).model_dump()
        assert output == {
            "headers": ["num", "prime"],
            "data": [[2, True], [3, True], [4, False]],
            "metadata": None,
        }
        with pytest.raises(ValueError):
            gr.Dataframe(type="unknown")  # type: ignore

        # When the headers don't match the data
        dataframe_output = gr.Dataframe(headers=["one", "two", "three"])
        output = dataframe_output.postprocess([[2, True], [3, True]]).model_dump()
        assert output == {
            "headers": ["one", "two"],
            "data": [[2, True], [3, True]],
            "metadata": None,
        }
        dataframe_output = gr.Dataframe(headers=["one", "two", "three"])
        output = dataframe_output.postprocess(
            [[2, True, "ab", 4], [3, True, "cd", 5]]
        ).model_dump()
        assert output == {
            "headers": ["one", "two", "three", "4"],
            "data": [[2, True, "ab", 4], [3, True, "cd", 5]],
            "metadata": None,
        }

        dataframe_output = gr.Dataframe(headers=["one", "two", "three"])
        output = dataframe_output.postprocess([(1, 2, 3), (4, 5, 6)]).model_dump()
        assert output == {
            "headers": ["one", "two", "three"],
            "data": [[1, 2, 3], [4, 5, 6]],
            "metadata": None,
        }

    def test_dataframe_postprocess_all_types(self):
        df = pd.DataFrame(
            {
                "date_1": pd.date_range("2021-01-01", periods=2),
                "date_2": pd.date_range("2022-02-15", periods=2).strftime(
                    "%B %d, %Y, %r"
                ),
                "number": np.array([0.2233, 0.57281]),
                "number_2": np.array([84, 23]).astype(np.int64),
                "bool": [True, False],
                "markdown": ["# Hello", "# Goodbye"],
            }
        )
        component = gr.Dataframe(
            datatype=["date", "date", "number", "number", "bool", "markdown"]  # type: ignore
        )
        output = component.postprocess(df).model_dump()
        assert output == {
            "headers": list(df.columns),
            "data": [
                [
                    pd.Timestamp("2021-01-01 00:00:00"),
                    "February 15, 2022, 12:00:00 AM",
                    0.2233,
                    84,
                    True,
                    "# Hello",
                ],
                [
                    pd.Timestamp("2021-01-02 00:00:00"),
                    "February 16, 2022, 12:00:00 AM",
                    0.57281,
                    23,
                    False,
                    "# Goodbye",
                ],
            ],
            "metadata": None,
        }

    def test_dataframe_postprocess_only_dates(self):
        df = pd.DataFrame(
            {
                "date_1": pd.date_range("2021-01-01", periods=2),
                "date_2": pd.date_range("2022-02-15", periods=2),
            }
        )
        component = gr.Dataframe(datatype=["date", "date"])  # type: ignore
        output = component.postprocess(df).model_dump()
        assert output == {
            "headers": list(df.columns),
            "data": [
                [
                    pd.Timestamp("2021-01-01 00:00:00"),
                    pd.Timestamp("2022-02-15 00:00:00"),
                ],
                [
                    pd.Timestamp("2021-01-02 00:00:00"),
                    pd.Timestamp("2022-02-16 00:00:00"),
                ],
            ],
            "metadata": None,
        }

    def test_dataframe_postprocess_styler(self):
        component = gr.Dataframe()
        df = pd.DataFrame(
            {
                "name": ["Adam", "Mike"] * 4,
                "gpa": [1.1, 1.12] * 4,
                "sat": [800, 800] * 4,
            }
        )
        s = df.style.format(precision=1, decimal=",")
        output = component.postprocess(s).model_dump()  # type: ignore
        assert output == {
            "data": [
                ["Adam", 1.1, 800],
                ["Mike", 1.12, 800],
                ["Adam", 1.1, 800],
                ["Mike", 1.12, 800],
                ["Adam", 1.1, 800],
                ["Mike", 1.12, 800],
                ["Adam", 1.1, 800],
                ["Mike", 1.12, 800],
            ],
            "headers": ["name", "gpa", "sat"],
            "metadata": {
                "display_value": [
                    ["Adam", "1,1", "800"],
                    ["Mike", "1,1", "800"],
                    ["Adam", "1,1", "800"],
                    ["Mike", "1,1", "800"],
                    ["Adam", "1,1", "800"],
                    ["Mike", "1,1", "800"],
                    ["Adam", "1,1", "800"],
                    ["Mike", "1,1", "800"],
                ],
                "styling": [
                    ["", "", ""],
                    ["", "", ""],
                    ["", "", ""],
                    ["", "", ""],
                    ["", "", ""],
                    ["", "", ""],
                    ["", "", ""],
                    ["", "", ""],
                ],
            },
        }

        df = pd.DataFrame(
            {
                "A": [14, 4, 5, 4, 1],
                "B": [5, 2, 54, 3, 2],
                "C": [20, 20, 7, 3, 8],
                "D": [14, 3, 6, 2, 6],
                "E": [23, 45, 64, 32, 23],
            }
        )

        t = df.style.highlight_max(color="lightgreen", axis=0)
        output = component.postprocess(t).model_dump()
        assert output == {
            "data": [
                [14, 5, 20, 14, 23],
                [4, 2, 20, 3, 45],
                [5, 54, 7, 6, 64],
                [4, 3, 3, 2, 32],
                [1, 2, 8, 6, 23],
            ],
            "headers": ["A", "B", "C", "D", "E"],
            "metadata": {
                "display_value": [
                    ["14", "5", "20", "14", "23"],
                    ["4", "2", "20", "3", "45"],
                    ["5", "54", "7", "6", "64"],
                    ["4", "3", "3", "2", "32"],
                    ["1", "2", "8", "6", "23"],
                ],
                "styling": [
                    [
                        "background-color: lightgreen",
                        "",
                        "background-color: lightgreen",
                        "background-color: lightgreen",
                        "",
                    ],
                    ["", "", "background-color: lightgreen", "", ""],
                    [
                        "",
                        "background-color: lightgreen",
                        "",
                        "",
                        "background-color: lightgreen",
                    ],
                    ["", "", "", "", ""],
                    ["", "", "", "", ""],
                ],
            },
        }

    def test_dataframe_hidden_columns(self):
        """Test that hidden columns are properly excluded from the output"""
        component = gr.Dataframe()
        df = pd.DataFrame(
            {"a": [1, 2, 3], "b": [4, 5, 6], "color": ["red", "blue", "green"]}
        )
        styled_df = df.style.hide(axis=1, subset=["color"])
        output = component.postprocess(styled_df).model_dump()
        assert output == {
            "data": [
                [1, 4],
                [2, 5],
                [3, 6],
            ],
            "headers": ["a", "b"],
            "metadata": {
                "display_value": [
                    ["1", "4"],
                    ["2", "5"],
                    ["3", "6"],
                ],
                "styling": [
                    ["", ""],
                    ["", ""],
                    ["", ""],
                ],
            },
        }

    def test_is_empty(self):
        """Test is_empty method with various data types"""
        df = gr.Dataframe()
        assert df.is_empty([])
        assert df.is_empty([[]])
        assert df.is_empty(np.array([]))
        assert df.is_empty(np.zeros((2, 0)))
        assert df.is_empty(None)
        assert df.is_empty({})
        assert df.is_empty({"data": [], "headers": ["a", "b"]})
        assert df.is_empty({"data": []})
        assert not df.is_empty({"data": [1, 2]})
        assert not df.is_empty([[1, 2], [3, 4]])
        assert not df.is_empty(pd.DataFrame({"a": [1, 2]}))
        assert not df.is_empty(pd.DataFrame({"a": [1, 2]}).style)

    def test_get_headers(self):
        """Test get_headers method with various data types"""
        df = gr.Dataframe()
        test_df = pd.DataFrame({"col1": [1, 2], "col2": [3, 4]})
        assert df.get_headers(test_df) == ["col1", "col2"]
        assert df.get_headers(test_df.style) == ["col1", "col2"]
        assert df.get_headers({"headers": ["a", "b"]}) == ["a", "b"]
        assert df.get_headers(np.array([[1, 2], [3, 4]])) == []
        assert df.get_headers(None) == []

    def test_get_cell_data(self):
        """Test get_cell_data method with various data types"""
        df = gr.Dataframe()
        test_data = [[1, 2], [3, 4]]
        test_df = pd.DataFrame({"col1": [1, 3], "col2": [2, 4]})
        assert df.get_cell_data(test_data) == [[1, 2], [3, 4]]
        assert df.get_cell_data(test_df) == [[1, 2], [3, 4]]
        assert df.get_cell_data({"data": test_data}) == [[1, 2], [3, 4]]

        styled_df = test_df.style
        styled_df.hide(axis=1, subset=["col2"])
        assert df.get_cell_data(styled_df) == [[1], [3]]

    def test_static_columns(self):
        # when static_columns is specified, it should be stored
        dataframe = gr.Dataframe(static_columns=[0, 1])
        assert dataframe.static_columns == [0, 1]

        # when static_columns is specified with column_count
        dataframe = gr.Dataframe(column_count=4, static_columns=[0, 1])
        assert dataframe.static_columns == [0, 1]
        assert dataframe.column_count == (4, "dynamic")

        # when static_columns is empty
        dataframe = gr.Dataframe(column_count=4, static_columns=[])
        assert dataframe.static_columns == []

        # when static_columns is None
        dataframe = gr.Dataframe(column_count=4, static_columns=None)
        assert dataframe.static_columns == []

        # when static_columns is not specified at all
        dataframe = gr.Dataframe(column_count=4)
        assert dataframe.static_columns == []

    def test_auto_datatype(self):
        df_headers = [
            "String",
            "Int",
            "Float",
            "Pandas Time",
            "Numpy Time",
            "Datetime",
            "Boolean",
        ]

        list_data = [
            [
                "Irish Red Fox",
                185000,
                4.2,
                pd.Timestamp("2017-01-01T12"),
                np.datetime64("now"),
                datetime(2022, 1, 1),
                True,
            ],
            [
                "Irish Badger",
                95000,
                8.5,
                pd.Timestamp("2018-01-01T12"),
                np.datetime64("now"),
                datetime(2023, 1, 1),
                True,
            ],
            [
                "Irish Otter",
                13500,
                5.5,
                pd.Timestamp("2025-01-01T12"),
                np.datetime64("now"),
                datetime(2024, 1, 1),
                False,
            ],
        ]
        np_data = np.array(list_data, dtype=object)
        pl_data = pl.DataFrame(list_data, schema=df_headers)

        pd_data = pd.DataFrame(list_data, columns=df_headers)  # type: ignore
        styler_data = pd_data.style.apply(
            lambda row: [
                "background-color: lightgreen" if row["Boolean"] else "" for _ in row
            ],
            axis=1,
        )

        result = ["str", "number", "number", "date", "date", "date", "bool"]

        dataframe = gr.Dataframe(
            value=pd_data, headers=df_headers, interactive=True, datatype="auto"
        )
        assert dataframe.datatype == result

        dataframe = gr.Dataframe(
            value=list_data, headers=df_headers, interactive=True, datatype="auto"
        )
        assert dataframe.datatype == result

        dataframe = gr.Dataframe(
            value=np_data, headers=df_headers, interactive=True, datatype="auto"
        )
        assert dataframe.datatype == result

        dataframe = gr.Dataframe(
            value=styler_data, headers=df_headers, interactive=True, datatype="auto"
        )
        assert dataframe.datatype == result

        dataframe = gr.Dataframe(
            value=pl_data, headers=df_headers, interactive=True, datatype="auto"
        )
        result = ["str", "number", "number", "date", "str", "date", "bool"]
        assert dataframe.datatype == result

        dataframe = gr.Dataframe(value=[], datatype="auto")
        assert dataframe.datatype == "str"

        dataframe = gr.Dataframe(value=[1, 2, 3], datatype="auto")
        assert dataframe.datatype == "str"

        dataframe = gr.Dataframe(value=np.array([]), datatype="auto")
        assert dataframe.datatype == "str"

        dataframe = gr.Dataframe(value=np.array([1, 2, 3]), datatype="auto")
        assert dataframe.datatype == "str"

        dataframe = gr.Dataframe(value=[[1, 2], [3, 4]], datatype="auto")
        assert dataframe.datatype == ["number", "number"]

        dataframe = gr.Dataframe(value=np.array([[1, 2], [3, 4]]), datatype="auto")
        assert dataframe.datatype == ["number", "number"]