python: how to convert if statements to enum or make it pythonic

Question

I don't quite know how to word this question. So here goes. I have this code which looks terrible. recommendation only takes 5 cases.

if avg_recommendation=='BUY':
    recommendation=5
if avg_recommendation=='OVERWEIGHT':
    recommendation=4
if avg_recommendation=='HOLD':
    recommendation=3
if avg_recommendation=='UNDERWEIGHT':
    recommendation=2
if avg_recommendation=='SELL':
    recommendation=1

And I want to make it pythonic. How do I do it? I've read about enum and it looks like it might be my solution. But I am open to any solution that is elegant to look at. I'm still using python2.7 Thanks.

Answer 1

Use a dict of key, values. Where key is the incoming state, and the value is the outgoing state. This is a highly Pythonic pattern and one you will likely reuse again and again. In many cases you will want to dispatch on the input state, where you should put functions as values in the dict.

recs = {
  'SELL':        1,
  'UNDERWEIGHT': 2,
  'HOLD':        3,
  'OVERWEIGHT':  4,
  'BUY':         5,
}

# this will fail correctly with a KeyError for non-expected states.
recommendation = recs[avg_recommendation]