Enumerations

I recently learned about Enums. Here's a fine ChatGPT example which is essentially the content of of Google search result. ChatGPT is like the old internet where you could read text on websites.

ChatGPT

Following is just verbatim output from the following request:

➜  ~ llm "give me an example of using Enums in python"

Sure! Here's an example of using Enums in Python:

from enum import Enum

# Define an Enum class
class Color(Enum):
    RED = 1
    GREEN = 2
    BLUE = 3

# Accessing Enum members
print(Color.RED)                    # Color.RED
print(Color.RED.name)               # RED
print(Color.RED.value)              # 1

# Iterating over Enum members
for color in Color:
    print(color)

# Checking Enum membership
print(Color.RED in Color)           # True
print(Color.YELLOW in Color)        # False

# Comparing Enum members
print(Color.RED == Color.RED)       # True
print(Color.RED == Color.GREEN)     # False

# Enum as dictionary keys
colors = {
    Color.RED: "FF0000",
    Color.GREEN: "00FF00",
    Color.BLUE: "0000FF"
}
print(colors[Color.RED])            # FF0000

In this example, we define an enum class called "Color" with three members - RED, GREEN, and BLUE. We can access enum members using dot notation (e.g., Color.RED). We can also access the name and value of an enum member using the name and value properties.

We can iterate over all the enum members using a for loop. Enum membership can be checked using the in keyword.

Enum members can be compared using the equality operator (==). We can also use enums as dictionary keys, as shown in the last example.

Back to me

This is all great. My takeaway is that we can replace strings with Enums across the codebase. The Color class enumerates color options and removes typos. There are two more things I want, however.

First, since this is an API interacting with other services, we need to be able to serialize to JSON. 

print(json.dumps({"inspection": Inspection.VISUAL}))
# TypeError: Object of type Inspection is not JSON serializable

print(json.dumps({"inspection": Inspection.VISUAL.value}))
{"inspection": "Visual Check"}

The default Enu, as you can see, is not serializable – but is value is. So all over the codebase we’d have to call .value whenever the Enum gets turned into JSON and that's not cool.

Well, can chatGPT help here? Maybe but in the time it'd take to fiddle with GPT I can do a Google search and get my answer on StackOverflow. The trick is to inherit from str and Enum. 

class Inspection(str, Enum):
   """JSON serializable inspection types with custom exception"""


   VISUAL = "Visual Check"
   BASIC = "Basic Inspection"
   STANDARD = "Standard Inspection"
   ADVANCED = "Advanced Inspection"


print(json.dumps({"inspection": Inspection.VISUAL}))
# {"inspection": "Visual Check"}

Boom! First problems solved and now the second problem.... A nice property of Enums is, like a dictionary, they raise an error if you try to make an Enum from an unknown value. For Example: 

print(Inspection("Just look it it real quick"))
# ValueError: 'Just look it it real quick' is not a valid Inspection

This is nice but we want to raise a custom exception (whether we should have custom exceptions is a topic for another day)

def test_unknown_inspection():
   with pytest.raises(MissingModelConfigError):
       Inspection("Give it a quick look")

No problem because as we can see in stacktrace

   575     @classmethod
    576     def _missing_(cls, value):
--> 577         raise ValueError("%r is not a valid %s" % (value, cls.__name__))

all we have to do is override the _missing_ method.

Result

Now we're done

class Inspection(str, Enum):
   """JSON serializable inspection types with custom exception"""


   VISUAL = "Visual Check"
   BASIC = "Basic Inspection"
   STANDARD = "Standard Inspection"
   ADVANCED = "Advanced Inspection"


   @classmethod
   def _missing_(cls, value):
       raise MissingModelConfigError(
           missing_property="inspection", expected_value=value
       )

For this effort we don’t have to worry about typos or an unknown inspection sneaking in. All we need to care about is a subset of these defined inspections. Very cool!

links

social