CODEWITHSUNDEEP
pythonpython-2.7enums
vt2424253
vt2424253

Reputation: 1427

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

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.

Upvotes: 3

Views: 8194

Answers (2)

abarnert
abarnert

Reputation: 365777

Using a dict is a reasonable way to make this Pythonic—but so is using an enum.

While the enum module wasn't built-in to the stdlib until Python 3.4, you can install enum34, a backport of that module that works on every version back to Python 2.4, or any of dozens of other enum packages if you like them better.

With enum34:

>>> import enum
>>> class Recommendation(enum.IntEnum):
...     BUY = 5
...     OVERWEIGHT = 4
...     HOLD = 3
...     UNDERWEIGHT = 2
...     SELL = 1
>>> Recommendation.BUY
<Recommendation.BUY: 5>
>>> Recommendation.BUY.value
5
>>> Recommendation(5)
<Recommendation.BUY: 5>
>>> recs = [Recommendation.BUY, Recommendation.HOLD, Recommendation.BUY]
>>> avg_rec = sum(recs)//len(rec)
>>> Recommendation(rec)
<Recommendation.OVERWEIGHT: 4>

(I assumed you wanted to do arithmetic on these things, given that you're using a variable named avg_recommendation. If not, you probably want to use Enum, not IntEnum.)

One advantage of this over dicts is the explicit repr, which makes things a lot easier to debug: you can see that you have a <Recommendation.OVERWEIGHT: 4>, not a 4 that may or may not be a recommendation.

Another is that you don't need to convert back and forth between the names and the values, which can be a major source of errors—especially in Python 2.x, where "SELL" == 1 is not an error, it's just false; Recommendation.SELL == 1 is true if you're using IntEnum, or an exception if you're using Enum, it never silently does the wrong thing.

Meanwhile, what if you have a string, like 'HOLD', and you want to get a Recommendation.HOLD out of that? Unfortunately, that's one of the features that was intentionally left out of the stdlib enum module to keep it simple.* (Read PEP 435 for a full discussion of all of the possible features that were left out and why.) So, if you try it, you'll get an error:

>>> Recommendation('HOLD')
ValueError: HOLD is not a valid Recommendation

But you could use one of the many other enum modules on PyPI. In particular, flufl.enum** works exactly like the stdlib in all the examples shown above, but also has a few extra features. So, if you pip install flufl.enum, then change the import enum to from flufl import enum, all of the above will work, and so will this:

>>> Recommendation('HOLD')
<EnumValue: Recommendation.HOLD [value=3]>

* If you only wanted string values, and not also numbers, then this would by easy, but obviously that won't work in your case.

** flufl.enum was one of the major influences to the stdlib module.

Upvotes: 3

Ali Afshar
Ali Afshar

Reputation: 41643

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]

Upvotes: 7

Related Questions