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
Enum
class from theenum
module. - We define an enum called
Sizes
with three members:SMALL
,MEDIUM
, andLARGE
. - In the
get_value_from_enum
function, we specify the parametersize
as typeSizes
. This tells the type checker that thesize
parameter must be one of the members from theSizes
enum. - Inside the function, we can access members of enum using the
.name
and.value
properties. - We also add the return type of the function, which is int as the
size.value
returns integer value. - The
get_value_from_enum
function is then called withSizes.MEDIUM
as 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
Literal
type from thetyping
module. - In
get_value_from_enum
function, thesize
parameter is type-hinted usingLiteral
, which specifies the accepted members ofSizes
. - The type checker will ensure that the
size
parameter 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.value
returns integer value. - The
get_value_from_enum
function is then called withSizes.MEDIUM
as 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
Literal
type: TheLiteral
approach 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.