Pandas style for loop not working correctly-CodePudding

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)

CodePudding user response：

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

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]
})

CodePudding user response：

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)