Pandas style for loop not working correctly

user_15325710 0
Viewed 383

I have a dataframe in Pandas containing planet names and information about each planet. Planet names are stored under the column with the header 'planet'. I want to highlight rows of specific planets, Earth, Saturn, and Mars. For some reason, when I run this script and export as an Excel file, only the first planet in my highlight_planets list is highlighted. In this case, only Earth is highlighted (Saturn and Mars are not highlighted). If I move Saturn into the first position and Earth into the second position, only Saturn rows will be highlighted in the dataframe. How do I highlight all rows where my planet column contains Earth, Saturn, or Mars?

Thanks.

def highlight_sentiment(row):

    highlight_planets = ['Earth', 'Saturn', 'Mars']

    for m in highlight_planets:


        if row['planet'] == m:
           return ['background-color: yellow'] * len(row)
        else:
           return ['background-color: white'] * len(row)

df.style.apply(highlight_sentiment, axis=1)

df = df.style.apply(highlight_sentiment, axis=1)
df.to_excel("df_test.xlsx", index = False)
pandas-styles pandas highlight python
2 Answers

In your for loop, you return immediately if the planet is not the first value in your list.

Move the last return statement outside the for loop so it only returns a white background if the planet doesn't match any value in your list.

def highlight_sentiment(row):
    highlight_planets = ['Earth', 'Saturn', 'Mars']
    for m in highlight_planets:
        if row['planet'] == m:
           return ['background-color: yellow'] * len(row)
    return ['background-color: white'] * len(row)

As an aside, you could avoid the for loop altogether and just use in:

def highlight_sentiment(row):
    highlight_planets = ['Earth', 'Saturn', 'Mars']
    if row['planet'] in highlight_planets:
        return ['background-color: yellow'] * len(row)
    return ['background-color: white'] * len(row)
user_98576310

The more standard approach to this problem is to build out a DataFrame of styles that can be built dynamically based on where planets are in the list using loc + Series.isin:

def highlight_sentiment(df_, highlight_planets):
    # Build DataFrame With Default Style
    styles_df = pd.DataFrame('background-color: white',
                             index=df_.index,
                             columns=df_.columns)
    # Update Styles based on condition
    styles_df.loc[
        df_['planet'].isin(highlight_planets), :
    ] = 'background-color: yellow'
    return styles_df


# Pass highlight_planets as an argument to allow more flexible styles
styler = df.style.apply(highlight_sentiment,
                        highlight_planets=['Earth', 'Saturn', 'Mars'],
                        axis=None)
# Use styler to write to excel
styler.to_excel("df_test.xlsx", index=False)

styled table

Sample Data and imports:

import pandas as pd

df = pd.DataFrame({
    'planet': ['Mercury', 'Venus', 'Earth',
               'Mars', 'Jupiter', 'Saturn'],
    'data': [1, 2, 3, 4, 5, 6]
})
user_154978880
Related