Search code examples
pythonpandasdataframepandas-styles

Adding color by string to a pandas dataframe


I'm working on a dataframe that I output to an html file. I would like to add colors to certain strings:

df1 = df1[(
    df1['Visitor'].str.contains('^TOR|^MTL|^CGY|^WPG|^VAN|^EDM|^OTT', na=False)
) | (df1['Home'].str.contains('^TOR|^MTL|^CGY|^WPG|^VAN|^EDM|^OTT', na=False))]

df1.fillna('', inplace=True)
df1.to_html('schedule.html', index=False)

The above code removes all rows that don't contain any of the listed teams. I would like to add colors to the rows that are left by string. for example 'TOR' would be colored blue.

As seen in another thread, this is what I tried but nothing changed:

    def styler(col):
  
    if col.name != 'Visitor vs Home':
        return [''] * len(col)

    bg_color = col.map({
        'TOR': 'blue',
        'MTL': 'red',
        'VAN': 'green',
    }).fillna('')  
    return 'background-color:' + bg_color


df1.style.apply(styler)

Any tips or suggestions are very much welcome.

Thanks!


Solution

  • Series.map can only be used for exact string matching. If needing to test for string contains, use str.contains with the appropriate patterns.

    We can then use np.select to allow mapping of pattern to colour. Lastly, we can take advantage of the subset parameter of Styler.apply to only style specific columns instead of having to check the name of each column passed in:

    def color_col(col):
        return np.select(
            [col.str.contains('TOR', na=False),
             col.str.contains('MTL', na=False),
             col.str.contains('VAN', na=False)],
            ['background-color: blue',
             'background-color: red',
             'background-color: green'],
            default=''
        ).astype(str)
    
    
    df.style.apply(color_col, subset=['Visitor vs Home'])
    

    We can also further parameterize our styling function so that we can pass a dictionary of patterns and corresponding colours instead:

    def color_col(col, pattern_map, default=''):
        return np.select(
            [col.str.contains(k, na=False) for k in pattern_map.keys()],
            [f'background-color: {v}' for v in pattern_map.values()],
            default=default
        ).astype(str)
    
    
    df.style.apply(color_col,
                   pattern_map={'TOR': 'blue',
                                'MTL': 'red',
                                'VAN': 'green'},
                   subset=['Visitor vs Home'])
    

    Both options produce:

    styled table


    Sample Data and imports:

    import numpy as np
    import pandas as pd
    
    df = pd.DataFrame({
        'Visitor': ['TOR', 'MTL', 'VAN', 'WPG', 'EDM'],
        'Home': ['CGY', 'WPG', 'OTT', 'TOR', 'MTL'],
        'Visitor vs Home': ['String with TOR',
                            'String with MTL',
                            'String with VAN',
                            'String with no match',
                            np.NaN]
    })