diff --git a/great_tables/_formats.py b/great_tables/_formats.py index 7b95767b0..5bfc6bc04 100644 --- a/great_tables/_formats.py +++ b/great_tables/_formats.py @@ -99,23 +99,25 @@ def fmt( The GT object is returned. This is the same object that the method is called on so that we can facilitate method chaining. """ - - # If a single function is supplied to `fns` then - # repackage that into a list as the `default` function - if isinstance(fns, Callable): - fns = FormatFns(default=fns) - row_res = resolve_rows_i(self, rows) row_pos = [name_pos[1] for name_pos in row_res] - col_res = resolve_cols_c(self, columns) - formatter = FormatInfo(fns, col_res, row_pos) + # If a single function is supplied to `fns` then + # repackage that into a list as the `default` function + formatter: list[FormatInfo] = [] + if isinstance(fns, Callable): + _fmtfs = FormatFns(default=fns) + formatter.append(FormatInfo(_fmtfs, col_res, row_pos)) + elif isinstance(fns, list): + for _fns in fns: + _fmtfs = FormatFns(default=_fns) + formatter.append(FormatInfo(_fmtfs, col_res, row_pos)) if is_substitution: - return self._replace(_substitutions=[*self._substitutions, formatter]) + return self._replace(_substitutions=[*self._substitutions, *formatter]) - return self._replace(_formats=[*self._formats, formatter]) + return self._replace(_formats=[*self._formats, *formatter]) def fmt_number( diff --git a/great_tables/_gt_data.py b/great_tables/_gt_data.py index 12919a3be..648ebbd3d 100644 --- a/great_tables/_gt_data.py +++ b/great_tables/_gt_data.py @@ -159,15 +159,17 @@ def render_formats(self, data_tbl: TblData, formats: list[FormatInfo], context: eval_func = getattr(fmt.func, context, fmt.func.default) if eval_func is None: raise Exception("Internal Error") + cell_info = [] for col, row in fmt.cells.resolve(): result = eval_func(_get_cell(data_tbl, row, col)) if isinstance(result, FormatterSkipElement): continue - - # TODO: I think that this is very inefficient with polars, so - # we could either accumulate results and set them per column, or - # could always use a pandas DataFrame inside Body? - _set_cell(self.body, row, col, result) + cell_info.append((row, col, result)) + # TODO: I think that this is very inefficient with polars, so + # we could either accumulate results and set them per column, or + # could always use a pandas DataFrame inside Body? + for _cell_info in cell_info: + data_tbl = _set_cell(self.body, *_cell_info) return self diff --git a/great_tables/_tbl_data.py b/great_tables/_tbl_data.py index b4722f3dd..04b5bb4e3 100644 --- a/great_tables/_tbl_data.py +++ b/great_tables/_tbl_data.py @@ -197,11 +197,13 @@ def _(data, row: int, column: str, value: Any) -> None: # if this is violated, get_loc will return a mask col_indx = data.columns.get_loc(column) data.iloc[row, col_indx] = value + return data @_set_cell.register(PlDataFrame) def _(data, row: int, column: str, value: Any) -> None: data[row, column] = value + return data # _get_column_dtype ---- @@ -297,6 +299,7 @@ def _( elif callable(expr): # TODO: currently, we call on each string, but we could be calling on # pd.DataFrame.columns instead (which would let us use pandas .str methods) + col_pos = {k: ii for ii, k in enumerate(list(data.columns))} return [(col, col_pos[col]) for col in data.columns if expr(col)] diff --git a/tests/test_formats.py b/tests/test_formats.py index afc1d4993..cf279f2e5 100644 --- a/tests/test_formats.py +++ b/tests/test_formats.py @@ -48,6 +48,22 @@ def test_format_fns(): assert res == ["2", "3"] +def test_format_multi_fns(): + df = pd.DataFrame({"x": [1, 2]}) + gt = GT(df) + new_gt = fmt(gt, fns=[lambda x: str(x + 1), lambda x: str(x + 2)], columns=["x"]) + + formats_fn = new_gt._formats[0] + + res = list(map(formats_fn.func.default, df["x"])) + assert res == ["2", "3"] + + formats_fn = new_gt._formats[1] + + res = list(map(formats_fn.func.default, df["x"])) + assert res == ["3", "4"] + + def test_format_snap(snapshot): new_gt = ( GT(exibble)