|
| 1 | +title: Making Enums (as always, arguably) more Pythonic |
| 2 | +author: Harry |
| 3 | +image: crab_nebula_multiple.png |
| 4 | +image_credit: https://commons.wikimedia.org/wiki/File:Crab_Nebula_in_Multiple_Wavelengths.png |
| 5 | + |
| 6 | +OK this isn't really anything to do with software architecture, but: |
| 7 | + |
| 8 | +> I hate [enums](https://docs.python.org/3/library/enum.html)! |
| 9 | +
|
| 10 | +I thought to myself, again and again, when having to deal with them recently. |
| 11 | + |
| 12 | +Why? |
| 13 | + |
| 14 | +```python |
| 15 | +class BRAIN(Enum): |
| 16 | + SMALL = 'small' |
| 17 | + MEDIUM = 'medium' |
| 18 | + GALAXY = 'galaxy' |
| 19 | +``` |
| 20 | + |
| 21 | +What could be wrong with that, I hear you ask? |
| 22 | +Well, accuse me of wanting to _stringly type_ everything if you will, |
| 23 | +but: those enums may look like strings but they aren't! |
| 24 | + |
| 25 | +```python |
| 26 | +assert BRAIN.SMALL == 'small' |
| 27 | +# nope, <BRAIN.SMALL: 'small'> != 'small' |
| 28 | + |
| 29 | +assert str(BRAIN.SMALL) == 'small' |
| 30 | +# nope, 'BRAIN.SMALL' != 'small' |
| 31 | + |
| 32 | +assert BRAIN.SMALL.value == 'small' |
| 33 | +# finally, yes. |
| 34 | +``` |
| 35 | + |
| 36 | +I imagine some people think this is a feature rather than a bug? But for me |
| 37 | +it's an endless source of annoyance. They look like strings! I defined them |
| 38 | +as strings! Why don't they behave like strings arg! |
| 39 | + |
| 40 | +Just one common motivating example: often what you want to do with those |
| 41 | +enums is dump them into a database column somewhere. This not-quite-a-string |
| 42 | +behaviour will cause your ORM or `db-api` library to complain like mad, and |
| 43 | +no end of footguns and headscratching when writing tests, custom SQL, and so on. |
| 44 | +At this point I'm wanting to throw them out and just use normal constants! |
| 45 | + |
| 46 | +But, one of the nice promises from Python's `enum` module is that **it's iterable**. |
| 47 | +So it's easy not just to refer to one constant, |
| 48 | +but also to refer to the list of all allowed constants. Maybe that's enough |
| 49 | +to want to rescue it? |
| 50 | + |
| 51 | +But, again, it doesn't quite work the way you might want it to: |
| 52 | + |
| 53 | +```python |
| 54 | +assert list(BRAIN) == ['small', 'medium', 'galaxy'] # nope |
| 55 | +assert [thing for thing in BRAIN] == ['small', 'medium', 'galaxy'] # nope |
| 56 | +assert [thing.value for thing in BRAIN] == ['small', 'medium', 'galaxy'] # yes |
| 57 | +``` |
| 58 | + |
| 59 | +Here's a _truly_ wtf one: |
| 60 | + |
| 61 | +```python |
| 62 | +assert random.choice(BRAIN) in ['small', 'medium', 'galaxy'] |
| 63 | +# Raises an Exception!!! |
| 64 | + |
| 65 | + File "/usr/local/lib/python3.9/random.py", line 346, in choice |
| 66 | + return seq[self._randbelow(len(seq))] |
| 67 | + File "/usr/local/lib/python3.9/enum.py", line 355, in __getitem__ |
| 68 | + return cls._member_map_[name] |
| 69 | +KeyError: 2 |
| 70 | +``` |
| 71 | + |
| 72 | +I have no idea what's going on there. What we actually wanted was |
| 73 | + |
| 74 | +```python |
| 75 | +assert random.choice(list(BRAIN)) in ['small', 'medium', 'galaxy'] |
| 76 | +# which is still not true, but at least it doesn't raise an exception |
| 77 | +``` |
| 78 | + |
| 79 | +Now the standard library does provide a solution |
| 80 | +if you want to duck-type your enums to integers, |
| 81 | +[IntEnum](https://docs.python.org/3/library/enum.html#derived-enumerations) |
| 82 | + |
| 83 | + |
| 84 | +```python |
| 85 | +class IBRAIN(IntEnum): |
| 86 | + SMALL = 1 |
| 87 | + MEDIUM = 2 |
| 88 | + GALAXY = 3 |
| 89 | + |
| 90 | +assert IBRAIN.SMALL == 1 |
| 91 | +assert int(IBRAIN.SMALL) == 1 |
| 92 | +assert IBRAIN.SMALL.value == 1 |
| 93 | +assert [thing for thing in IBRAIN] == [1, 2, 3] |
| 94 | +assert list(IBRAIN) == [1, 2, 3] |
| 95 | +assert [thing.value for thing in IBRAIN] == [1, 2, 3] |
| 96 | +assert random.choice(IBRAIN) in [1, 2, 3] # this still errors but: |
| 97 | +assert random.choice(list(IBRAIN)) in [1, 2, 3] # this is ok |
| 98 | +``` |
| 99 | + |
| 100 | +That's all fine and good, but I don't _want_ to use integers. |
| 101 | +I want to use strings, because then when I look in my database, |
| 102 | +or in printouts, or wherever, |
| 103 | +the values will make sense. |
| 104 | + |
| 105 | +Well, the [docs say](https://docs.python.org/3/library/enum.html#others) |
| 106 | +you can just subclass `str` and make your own `StringEnum` that will work just like `IntEnum`. |
| 107 | +But it's LIES: |
| 108 | + |
| 109 | +```python |
| 110 | +class BRAIN(str, Enum): |
| 111 | + SMALL = 'small' |
| 112 | + MEDIUM = 'medium' |
| 113 | + GALAXY = 'galaxy' |
| 114 | + |
| 115 | +assert BRAIN.SMALL.value == 'small' # ok, as before |
| 116 | +assert BRAIN.SMALL == 'small' # yep |
| 117 | +assert list(BRAIN) == ['small', 'medium', 'galaxy'] # hooray! |
| 118 | +assert [thing for thing in BRAIN] == ['small', 'medium', 'galaxy'] # hooray! |
| 119 | +random.choice(BRAIN) # this still errors but ok i'm getting over it. |
| 120 | + |
| 121 | +# but: |
| 122 | +assert str(BRAIN.SMALL) == 'small' #NOO!O!O! 'BRAIN.SMALL' != 'small' |
| 123 | +# so, while BRAIN.SMALL == 'small', str(BRAIN.SMALL) != 'small' aaaargh |
| 124 | +``` |
| 125 | + |
| 126 | +So here's what I ended up with: |
| 127 | + |
| 128 | +```python |
| 129 | +class BRAIN(str, Enum): |
| 130 | + SMALL = 'small' |
| 131 | + MEDIUM = 'medium' |
| 132 | + GALAXY = 'galaxy' |
| 133 | + |
| 134 | + def __str__(self) -> str: |
| 135 | + return str.__str__(self) |
| 136 | +``` |
| 137 | + |
| 138 | +* this basically avoids the need to use `.value` anywhere at all in your code |
| 139 | +* enum values duck type to strings in the ways you'd expect |
| 140 | +* you can iterate over brain and get string-likes out |
| 141 | + - altho `random.choice()` is still broken, i leave that as an exercise for the reader |
| 142 | +* and type hints still work! |
| 143 | + |
| 144 | +```python |
| 145 | +# both of these type check ok |
| 146 | +foo = BRAIN.SMALL # type: str |
| 147 | +bar = BRAIN.SMALL # type: BRAIN |
| 148 | +``` |
| 149 | + |
| 150 | +Example code is [in a Gist](https://gist.github.com/hjwp/405f04802ea558f042728ec5edbb4e62) |
| 151 | +if you want to play around. |
| 152 | +Let me know if you find anything better! |
0 commit comments