Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

How to add type hint for dynamically created enums?

Consider than I'm creating an Enum class as below,

from enum import Enum

key_value = {"DAY": "d", "WEEK": "w", "MONTH": "m"}
FooEnum = Enum("FooEnum", key_value)

How can make the FooEnum class represent proper type hints (to identify its attributes, such as DAY, WEEK and MONTH)?

Note: In real cases, the size of key_value will be a bit large - say 100 entries.

like image 818
JPG Avatar asked Mar 25 '26 06:03

JPG


1 Answers

Indeed - if your enum is dynamically built, it is "data" rather than code. Static type checking in Python deals with what is code. If you are writing "myvar.DAY" in your code, and DAY is an enumeration member created by a string obtained as data, Python static checking can't cope with it: they are from "different universes".

The most straight foward thing to do is simply ignore type hinting on the lines that hardcode member names for a dynamically created Enum.

(fun fact, that you probably hit when trying thi: mypy will error on the Enum call in your snippet: it expects only literal dictionaries typed inline in the call to Enum, and will refuse to introspect the variable to assert it is a dictionary. In other words, the line FooEnum = Enum("FooEnum", key_value) makes mypy error with " error: Enum() expects a string, tuple, list or dict literal as the second argument")

So, the dilema is set: mypy does static code checking. Somehow maneuvering to dynamically create a protocol (a subclass of typing.Protocol) out of the dynamic data and registering your Enum as an implementation of that can't work either - mypy won't be able to static check for the protocol validity.

What may be possible, apart from marking all lines consuming the Enum with # type: ignore as I stated above, is creating a protocol with a static subset of your enum members - and hardcode in your ".py" file at least the members that are explicitly used in code.

In "businness as usual", mypy would automatically identify classes that cope with a protocol, with no changes needed. This is not the case here: mypy can't know what is inside the enum anyway, so we'd need to explicitly call the register method on the protocol to indicate each Enum that implements it. Now, there is another problem: due to the "promiscuous" relation Enum values and their classes have, you can't just register an Enum subclass as implementing a protocol and expect it to work - mypy will check for an instance of your enumeration, and will "see" its class being passed, and error.

So, not only one has to hardcode the subset of attributes that will actually be used in code, one also needs to explicitly cast (typing.cast) the enumeration to the protocol class before calling the place the Enum will be used.

import typing as t
import json
from enum import Enum

key_value: dict = {"DAY": "d", "WEEK": "w", "MONTH": "m"}
DateEnum = Enum("DateEnum", key_value) # type: ignore
# ^  you can't escape from this "ignore" anyway.
TimeStampEnum = Enum("TimeStampEnum", json.load(open("file_with_time_units.json"))) # type: ignore


class MinimalTimeUnits(t.Protocol):
    # we only care about using these in code:
    DAY: str
    MONTH: str


def blah(val: MinimalTimeUnits)-> None:
    # Code uses just the enumeration members defined in the procotol:
    x = f"{val.DAY} of {val.MONTH}"


some_condition = True  # <- here so mypy won't error on missing variable on this example

def somecode()->None:
    ...
    time_enum = DateEnum if some_condition else TimeStampEnum
    blah(t.cast(MinimalTimeUnits, time_enum))
like image 98
jsbueno Avatar answered Mar 27 '26 18:03

jsbueno