Type Hint Enum in Python
Last Updated :
09 Feb, 2024
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
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}" )
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}" )
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}" )
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}" )
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.
Share your thoughts in the comments
Please Login to comment...