Type hinting enum member value in Python

Question

For an enum Foo, How could one type hint a variable that must contain a member value (not the member itself) of some Enum-- e.g. a value such that Foo(x) will return a valid member of Foo?

Here's a simplified version of my motivating example:

class DisbursementType(Enum):
    DISBURSEMENT = "disbursement"
    REFUND = "refund"
    ROLLBACK = "rollback"

class SerializedDisbursement(TypedDict):
    transaction_type: ???
    id: str
    amount: float

a: SerializedDisbursement = {"transaction_type": "refund", id: 1, amount: 4400.24}

I would really like to avoid simply typeing transaction_type as Literal['disbursement', 'refund', 'rollback'] as that would be quite prone to getting out of synch over time.

Answer 1

The most widely compatible option is to just have an assertion that validates that the literal type doesn't go out of sync with the enum values:

class DisbursementType(enum.Enum):
    DISBURSEMENT = "disbursement"
    REFUND = "refund"
    ROLLBACK = "rollback"

DisbursementValue = typing.Literal['disbursement', 'refund', 'rollback']

assert set(typing.get_args(DisbursementValue)) == {member.value for member in DisbursementType}

class SerializedDisbursement(typing.TypedDict):
    transaction_type: DisbursementValue
    id: str
    amount: float

This ensures maximum compatibility with static analyzers, but requires repeating all member values. Also, the assertion cannot be checked statically.

---

Other options break static analysis. For example, if you use the functional API to create the enum from the literal type:

DisbursementValue = typing.Literal['disbursement', 'refund', 'rollback']

DisbursementType = enum.Enum('DisbursementType',
                             {name.upper(): name for name in typing.get_args(DisbursementValue)})

then mypy doesn't understand the enum, and at that point, there's little point having annotations at all.

Similarly, if you try to use non-literal type arguments for the Literal type:

DisbursementValue = typing.Literal[tuple(member.value for member in DisbursementType)]

then that breaks too.

Answer 2

While this doesn't answer the question as asked, I think you might also reconsider why you're using a TypedDict with a string instead of a proper class to hold the enum (instead of a str, since DisbursementType really does seem like an enum), and which can then employ some custom serialization logic.

For example:

import dataclasses as dc
import json
from enum import Enum


class Transaction(Enum):
    DISBURSEMENT = "disbursement"
    REFUND = "refund"
    ROLLBACK = "rollback"

    def __str__(self):
        return self.value


@dc.dataclass
class Disbursement:
    transaction: Transaction
    id_: str
    amount: float

    def __str__(self):
        return json.dumps(dc.asdict(self), default=str)


if __name__ == "__main__":
    disbursement = Disbursement(
        Transaction.REFUND,
        "1",
        4400.24,
    )
    print(disbursement)

$ mypy example.py
Success: no issues found in 1 source file
$ python3 example.py
{"transaction": "refund", "id_": "1", "amount": 4400.24}

Alternatively, you can have your enum inherit from str and simplify a few things:

import dataclasses as dc
import json
from enum import Enum


class Transaction(str, Enum):
    DISBURSEMENT = "disbursement"
    REFUND = "refund"
    ROLLBACK = "rollback"


@dc.dataclass
class Disbursement:
    transaction: Transaction
    id_: str
    amount: float

    def __str__(self):
        return json.dumps(dc.asdict(self))


if __name__ == "__main__":
    disbursement = Disbursement(
        Transaction.REFUND,
        "1",
        4400.24,
    )
    print(disbursement)

Other considerations:

• Python 3.11 will have StrEnum: https://docs.python.org/3.11/howto/enum.html#strenum
• Notes on subclassing enum: https://docs.python.org/3/library/enum.html#others
• Equality comparisons with strings should probably be avoided: https://www.cosmicpython.com/blog/2020-10-27-i-hate-enums.html (prefer Disbursement.transaction == Transaction.REFUND to Disbursement.transaction == "refund" or str(Disbursement.transaction) == "refund"
• Inheriting from (str, Enum) may have some drawbacks: How to deserialise enumeration with string representation?

Finally, I wanted to note that defining __str__ on my Enum did not do what I expected it to do using your TypedDict example above; that's because str(mydict) calls repr to provide each of mydict.values:

class Example:
    def __repr__(self):
        print("I called repr!")
        return "from repr"

    def __str__(self):
        print("I called str!")
        return "from str"


if __name__ == "__main__":
    print(f"example: {Example()}\n")

    d = {"example": Example()}
    print(f"in a dict: {d}")

$ python3 foo.py
I called str!
example: from str

I called repr!
in a dict: {'example': from repr}

Additionally, you can't add custom methods to a TypedDict; if you change Example to inherit from TypedDict and rerun that last example, you'll see that neither __repr__ nor __str__ is called, and unfortunately there is no runtime error either (mypy helpfully warns error: Invalid statement in TypedDict definition; expected "field_name: field_type"). Because serialization logic seems to belong to Disbursement, I changed it to a somewhat similar class that allows me to customize its __str__: a dataclass.