_utils.py 585 Bytes
Newer Older
Philip Meier's avatar
Philip Meier committed
1
import enum
2
3
4
from typing import TypeVar, Type

T = TypeVar("T", bound=enum.Enum)
Philip Meier's avatar
Philip Meier committed
5
6
7
8
9


class StrEnumMeta(enum.EnumMeta):
    auto = enum.auto

10
    def from_str(self: Type[T], member: str) -> T:  # type: ignore[misc]
Philip Meier's avatar
Philip Meier committed
11
12
13
14
15
16
17
18
19
20
        try:
            return self[member]
        except KeyError:
            # TODO: use `add_suggestion` from torchvision.prototype.utils._internal to improve the error message as
            #  soon as it is migrated.
            raise ValueError(f"Unknown value '{member}' for {self.__name__}.") from None


class StrEnum(enum.Enum, metaclass=StrEnumMeta):
    pass