posit-dev / great-tables

Make awesome display tables using Python.
https://posit-dev.github.io/great-tables/
MIT License
1.43k stars 48 forks source link

Modify `CellStyle` to accommodate `pandas.Series.case_when` #280

Closed jrycw closed 2 months ago

jrycw commented 2 months ago

Hello team,

I noticed that with the introduction of pandas.Series.case_when in pandas 2.2, there's a similarity to polars.when-then-otherwise used in our tutorial. After investigating further, I discovered that including pdSeries in the isinstance check of the CellStyle._evaluate_expressions method might offer a solution. Here's the proposed code snippet:

@dataclass
class CellStyle:

    def _evaluate_expressions(self, data: TblData) -> Self:
             ...
            if isinstance(attr, (PdSeries, PlExpr)) or callable(attr):
                ...

This change allows us to use the following code:

def get_temp_color(df_):
    col = df_["Temp"]
    caselist = [(col.gt(70), "lightyellow"), (~col.gt(70), "lightblue")]
    return col.case_when(caselist)

gt_air.tab_style(
    style=style.fill(color=get_temp_color),
    locations=loc.body("Temp")
)

Initial testing indicates that the proposed solution works as expected:

def test_cell_value_from_pandas_series():
    df = pd.DataFrame({"x": [1, 2]})

    def get_color(df_):
        col = df_["x"]
        caselist = [(col.lt(2), "RED"), (~col.lt(2), "BLUE")]
        return col.case_when(caselist)

    style = CellStyleText(color=get_color)

    new_style = style._evaluate_expressions(df)._from_row(df, 0)

    assert new_style.color == "RED"

I'm interested in hearing your thoughts on this approach. Do you believe it's a viable idea, or do you foresee any potential issues with existing code? If we reach a consensus that this is a beneficial change, I'm prepared to submit a PR.

machow commented 2 months ago

Hey, thanks for raising this. I tested your code, and it seems to work on main?

image

I think this behavior is intentional, and is documented in Styling the Table Body -> Using Functions in the guide.

In any event, I think this approach with pandas makes sense, and hopefully works as is in Great Tables. If anything is needed to make it work, very happy to fix!

jrycw commented 2 months ago

Thank you for your response. Initially, I was exploring the possibility of not creating a background in the Combining Styling with cols_hide and didn't realize the callable option you mentioned. Thank you once again for your help.