python：如何将if语句转换为枚举或更具Python风格

Question

我不太知道该怎么问这个问题。好吧，我试试。我有一段代码，看起来很糟糕。recommendation 只处理5种情况。

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

我想让它变得更符合Python的风格。该怎么做呢？我读过关于枚举（enum）的内容，感觉这可能是我的解决方案。不过，我也愿意接受任何看起来优雅的解决办法。我现在还在用Python 2.7。谢谢。

Answer 1

使用字典是一种合理的方式来让代码更符合Python的风格，但使用枚举也是一个不错的选择。

虽然从Python 3.4开始，enum模块才被包含在标准库中，但你可以安装enum34，这是一个可以在Python 2.4及以上版本上使用的旧版本模块，或者你也可以选择其他许多枚举包，如果你觉得它们更好用。

使用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>

（我假设你想对这些值进行数学运算，因为你使用了一个叫avg_recommendation的变量。如果不是的话，你可能想用Enum而不是IntEnum。）

使用枚举的一个好处是它有明确的repr，这让调试变得更简单：你可以看到你有一个<Recommendation.OVERWEIGHT: 4>，而不是一个可能是推荐的4。

另一个好处是你不需要在名称和数值之间来回转换，这可以避免很多错误，尤其是在Python 2.x中，"SELL" == 1并不是错误，只是返回假；而Recommendation.SELL == 1在使用IntEnum时为真，使用Enum时会抛出异常，它不会默默地做错事。

---

那么，如果你有一个字符串，比如'HOLD'，想要得到Recommendation.HOLD，该怎么办呢？不幸的是，这是标准库enum模块故意省略的一个功能，以保持简单性。（想了解更多，可以阅读PEP 435，里面详细讨论了所有被省略的功能及原因。）所以，如果你尝试这样做，你会遇到错误：

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

不过你可以使用PyPI上许多其他的枚举模块。特别是，flufl.enum**的工作方式和标准库中的枚举模块完全一样，但它还增加了一些额外的功能。所以，如果你pip install flufl.enum，然后把import enum改成from flufl import enum，那么上面的所有内容都能正常工作，包括这个：

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

---

_{* 如果你只想要字符串值，而不需要数字，那会很简单，但显然这不适合你的情况。}

_{** flufl.enum对标准库模块的设计影响很大。}

Answer 2

使用一个字典，里面存放键和值。键是输入的状态，而值是输出的状态。这是一种非常符合Python风格的写法，你可能会多次用到。在很多情况下，你会想根据输入的状态来处理事情，这时候可以把函数放在字典的值里。

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]