Type Hint Enum in Python

Type hinting in Python has become a crucial aspect of writing clean and maintainable code. Enums, short for enumerations, are a convenient way to represent a set of named values. Combining type hinting with enums enhances code readability and makes it more robust. In this article, we will explore some simple and commonly used methods to type hint enums in Python.

How to Type Hint Enum in Python?

Below, are the methods of How to Type Hint Enum in Python.

• Basic Type Hinting
• Union Type Hinting
• Enum IntEnum Type Hinting
• Enum Member Type Hinting

Basic Type Hinting

The most straightforward method is to use basic type hinting syntax. For example, if you have an enum named Color, you can type hint a variable as follows: Here, the color parameter in the print_color the function is type hinted with the Color enum. This helps catch potential errors and improves code clarity.

Python3

from enum import Enum

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

def print_color(color: Color) -> None:
    print(f"Selected color: {color.name}")

# Example usage
selected_color: Color = Color.RED
print_color(selected_color)

Output

Selected color: RED

Type Hint Enum in Python Union Type Hinting

Another commonly used method involves using the Union type hint to allow multiple enum types or other types. In this example, the print_item function can accept either a Shape or Size enum, providing flexibility in the usage.

Python3

from enum import Enum
from typing import Union

class Shape(Enum):
    CIRCLE = 1
    SQUARE = 2

class Size(Enum):
    SMALL = 1
    MEDIUM = 2

def print_item(item: Union[Shape, Size]) -> None:
    print(f"Selected item: {item.name}")

# Example usage
selected_shape: Shape = Shape.CIRCLE
print_item(selected_shape)

selected_size: Size = Size.MEDIUM
print_item(selected_size)

Output

Selected item: CIRCLE
Selected item: MEDIUM

Type Hint Enum in Python Enum IntEnum Type Hinting

If you are using IntEnum for your enums, you can utilize the Enum type hint from the enum module. Here, the Priority enum is an IntEnum, and the priority parameter in the set_priority function is type hinted with Priority.

Python3

from enum import Enum, IntEnum

class Priority(IntEnum):
    LOW = 1
    MEDIUM = 2
    HIGH = 3

def set_priority(priority: Priority) -> None:
    print(f"Set priority to: {priority.name}")

# Example usage
selected_priority: Priority = Priority.MEDIUM
set_priority(selected_priority)

Output

Set priority to: MEDIUM

Type Hint Enum in Python Enum Member Type Hinting

For scenarios where you want to specify the type of an individual enum member, you can use the typing module’s Type class. In this example, the Type[Animal] type hint is applied to the animal_type parameter, explicitly stating that it should be an enum member of the Animal enum.

Python3

from enum import Enum
from typing import Type

class Animal(Enum):
    DOG = 'dog'
    CAT = 'cat'

def print_animal_type(animal_type: Type[Animal]) -> None:
    print(f"Selected animal type: {animal_type.name}")

# Example usage
selected_animal_type: Type[Animal] = Animal.DOG
print_animal_type(selected_animal_type)

Output

Selected animal type: DOG

Conclusion

Type hinting enums in Python contributes to code readability and helps catch potential errors during development. By using these five simple and commonly used methods, you can make your code more robust and maintainable while leveraging the power of enums. Whether using basic type hinting, union types, or more specific enum type hints, choose the approach that best fits your project’s requirements.