Python pandas sort dataframe by enum class values

user_3225309 0
Viewed 4497

If I have enum class:

from enum import Enum
class Colors(Enum):
    RED = 1
    ORANGE = 2
    GREEN = 3

And if I have a dataframe whose one column is color (it can be in lowercase to):

>>> import pandas as pd
>>> df = pd.DataFrame({'X':['A', 'B', 'C', 'A'], 'color' : ['GREEN', 'RED', 'ORANGE', 'ORANGE']})
>>> df
   X   color
0  A   GREEN
1  B     RED
2  C  ORANGE
3  A  ORANGE

How to make color column as categorical type respecting Color class values, and sort the dataframe by "color" and "X" (ascending)?

For example, the dataframe above should be sorted as:

X, color
--------
B, RED
A, ORANGE
C, ORANGE
A, GREEN
pandas categorical-data enums python-3.x sorting
2 Answers

Combination of this answer and this one: use a pd.Categorical to sort by the Colors class (with a slight edit to change its str):

from enum import Enum
import pandas as pd

df = pd.DataFrame({'X':['A', 'B', 'C', 'A'], 'color' : ['GREEN', 'RED', 'ORANGE', 'ORANGE']})

class Colors(Enum):
    RED = 1
    ORANGE = 2
    GREEN = 3
    def __str__(self):
        return self.name

df['color'] = pd.Categorical(df['color'], [str(i) for i in Colors], ordered=True)
df = df.sort_values(['color','X'])

Result:

   X   color
1  B     RED
3  A  ORANGE
2  C  ORANGE
0  A   GREEN
user_133869790

Use getattr:

df["value"] = df["color"].apply(lambda x: getattr(Colors, x).value)
df.sort_values(by=['value',"X"])

Output:

    X   color   value
1   B   RED     1
3   A   ORANGE  2
2   C   ORANGE  2
0   A   GREEN   3

In one line (and without creation of value column): 

df.iloc[pd.concat([df["X"], df["color"].apply(lambda x: getattr(Colors, x))], axis=1).sort_values(by=['color',"X"]).index]

Output:

    X   color
1   B   RED
3   A   ORANGE
2   C   ORANGE
0   A   GREEN
user_133280100
Related