Skip to main content

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 the enum module.
  • We define an enum called Sizes with three members: SMALL, MEDIUM, and LARGE.
  • In the get_value_from_enum function, we specify the parameter size as type Sizes. This tells the type checker that the size parameter must be one of the members from the Sizes 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 with Sizes.MEDIUM as its argument.
note

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 the typing module.
  • In get_value_from_enum function, the size parameter is type-hinted using Literal, which specifies the accepted members of Sizes.
  • The type checker will ensure that the size parameter is one of these literal enum members: Sizes.SMALL, Sizes.MEDIUM, or Sizes.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 with Sizes.MEDIUM as its argument.
note

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: The Literal 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.

note

In most cases, using the enum class directly is the recommended and cleaner way to type hint enums.