CODEWITHSUNDEEP
pythonenumspython-typing
Nentuaby
Nentuaby

Reputation: 600

Type hinting enum member value in Python

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.

Upvotes: 12

Views: 12813

Answers (2)

n8henrie
n8henrie

Reputation: 2975

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:

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.

Upvotes: 3

user2357112
user2357112

Reputation: 281958

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.

Upvotes: 13

Related Questions