ENH: Styler.apply with axis=None can accept ndarray (original) (raw)

Currently Styler.apply with axis=0 and axis=1 permit functions which return array-like objects, provided they are of appropriate shape. With axis=None the return must be a DataFrame.

This means the following example:

def highlight_max(d):
    """Highlight the maximum in a Series or DataFrame bold-orange."""
    css = 'background-color: orange; font-weight: bold;'
    return np.where(d == np.nanmax(d.values), css, None)

df.style.apply(highlight_max, axis=0)  # works
df.style.apply(highlight_max, axis=1)  # works
df.style.apply(highlight_max, axis=None)  # raises TypeError since return is not DataFrame

Describe the solution you'd like

The solution is a relatively simple change to code as part of the Styler._apply method that checks for array-like (or specifically ndarray) result instead of DataFrame, and constructs a DataFrame from that data if needs be, with appropriate index and cols.

API breaking implications

None since all previous behaviour is replicable.