How to Type Hint Enums in Python
Enums (enumerations) provide a way to create sets of named constants in Python, making your code more readable and maintainable. When working with enums, using type hints can further enhance code quality by allowing type checkers to catch errors early.
Type Hinting Enums with the Enum Class
The most straightforward and recommended way to type hint an enum is by using the enumeration class itself as the type annotation.
from enum import Enum
class Sizes(Enum):
SMALL = 1
MEDIUM = 2
LARGE = 3
def get_value_from_enum(size: Sizes) -> int:
print(size.name)
print(size.value)
return size.value
result = get_value_from_enum(Sizes.MEDIUM)
print(result)
Explanation:
- We import the
Enumclass from theenummodule. - We define an enum called
Sizeswith three members:SMALL,MEDIUM, andLARGE. - In the
get_value_from_enumfunction, we specify the parametersizeas typeSizes. This tells the type checker that thesizeparameter must be one of the members from theSizesenum. - Inside the function, we can access members of enum using the
.nameand.valueproperties. - We also add the return type of the function, which is int as the
size.valuereturns integer value. - The
get_value_from_enumfunction is then called withSizes.MEDIUMas its argument.
By using the enum class as a type hint, the type checker will verify that the function arguments are valid enum members, catching potential errors if a value outside the enum is used.
Type Hinting Enums with the Literal Type
Another approach is to use the Literal type from the typing module. The Literal type allows you to specify that a variable or parameter must have one of several specific literal values, like the enum members.
from typing import Literal
from enum import Enum
class Sizes(Enum):
SMALL = 1
MEDIUM = 2
LARGE = 3
def get_value_from_enum(size: Literal[Sizes.SMALL, Sizes.MEDIUM, Sizes.LARGE]) -> int:
print(size.name)
print(size.value)
return size.value
result = get_value_from_enum(Sizes.MEDIUM)
print(result)
Explanation:
- We import the
Literaltype from thetypingmodule. - In
get_value_from_enumfunction, thesizeparameter is type-hinted usingLiteral, which specifies the accepted members ofSizes. - The type checker will ensure that the
sizeparameter is one of these literal enum members:Sizes.SMALL,Sizes.MEDIUM, orSizes.LARGE. - We also add the return type of the function, which is int as the
size.valuereturns integer value. - The
get_value_from_enumfunction is then called withSizes.MEDIUMas its argument.
This method works well for smaller enums and is particularly useful if you want to restrict the function to a subset of the enum members.
Choosing the Right Approach
-
Using the Enum class directly: This approach is generally preferred for its simplicity and readability. It's concise and clearly indicates that the parameter should be an instance of that specific enum.
-
Using the
Literaltype: TheLiteralapproach can be useful when you need to explicitly list out the accepted values, and you want type checkers to enforce that only a specific combination of values are used in your function, not all the members of the enum. However, it can become verbose for larger enums.
In most cases, using the enum class directly is the recommended and cleaner way to type hint enums.