|
3 | 3 | import multiprocessing |
4 | 4 | from enum import IntEnum |
5 | 5 | from multiprocessing.synchronize import Lock, RLock |
6 | | -from typing import Any, Generic, TypeVar, get_args |
| 6 | +from types import get_original_bases |
| 7 | +from typing import Any, get_args, get_origin |
7 | 8 |
|
8 | 9 | from ._version import version as __version__ # noqa: F401 |
9 | 10 |
|
10 | | -T = TypeVar("T", bound=IntEnum) |
11 | 11 |
|
12 | | - |
13 | | -class IntEnumValue(Generic[T]): |
| 12 | +class IntEnumValue[T: IntEnum]: |
14 | 13 | """A multiprocessing safe shared object for `IntEnum` enum values.""" |
15 | 14 |
|
16 | 15 | # Pre-initialize type here to avoid numerous type ignores elsewhere |
17 | 16 | EnumType: type[T | IntEnum] = IntEnum |
18 | 17 |
|
19 | 18 | @classmethod |
20 | 19 | def __init_subclass__(cls) -> None: # noqa: D105 |
| 20 | + # find our base class among potentially multiple bases |
| 21 | + orig_base = None |
| 22 | + for base in get_original_bases(cls): |
| 23 | + origin = get_origin(base) |
| 24 | + if origin is not None and issubclass(origin, IntEnumValue): |
| 25 | + orig_base = base |
| 26 | + if orig_base is None: |
| 27 | + message = "Can not determine base class!" |
| 28 | + raise RuntimeError(message) |
| 29 | + |
21 | 30 | # set EnumType to the specific type specified by subclass |
22 | | - orig_base = cls.__orig_bases__[0] # type: ignore[attr-defined] |
23 | 31 | cls.EnumType = get_args(orig_base)[0] |
24 | 32 |
|
25 | 33 | def __init__(self, value: T | str, lock: None | Lock | RLock = None) -> None: |
|
0 commit comments