1
|
|
|
import enum |
|
|
|
|
2
|
|
|
from typing import Optional, Union |
3
|
|
|
|
4
|
|
|
from mandos import logger |
5
|
|
|
|
6
|
|
|
|
7
|
|
|
class TrueFalseUnknown(enum.Enum): |
|
|
|
|
8
|
|
|
true = enum.auto() |
9
|
|
|
false = enum.auto() |
10
|
|
|
unknown = enum.auto() |
11
|
|
|
|
12
|
|
|
@classmethod |
13
|
|
|
def parse(cls, s: str): |
|
|
|
|
14
|
|
|
tf_map = { |
15
|
|
|
"t": TrueFalseUnknown.true, |
16
|
|
|
"f": TrueFalseUnknown.false, |
17
|
|
|
"true": TrueFalseUnknown.true, |
18
|
|
|
"false": TrueFalseUnknown.false, |
19
|
|
|
} |
20
|
|
|
return tf_map.get(s.lower().strip(), TrueFalseUnknown.unknown) |
21
|
|
|
|
22
|
|
|
|
23
|
|
|
class CleverEnum(enum.Enum): |
24
|
|
|
""" |
25
|
|
|
An enum with a ``.of`` method that finds values |
26
|
|
|
with limited string/value fixing. |
27
|
|
|
May support an "unmatched" type -- a fallback value when there is no match. |
28
|
|
|
This is similar to pocketutils' simpler ``SmartEnum``. |
29
|
|
|
It is mainly useful for enums corresponding to concepts in ChEMBL and PubChem, |
30
|
|
|
where it's acceptable for the user to input spaces (like the database concepts use) |
31
|
|
|
rather than the underscores that Python requires. |
32
|
|
|
""" |
33
|
|
|
|
34
|
|
|
@classmethod |
35
|
|
|
def _unmatched_type(cls) -> Optional[__qualname__]: |
36
|
|
|
return None |
37
|
|
|
|
38
|
|
|
@classmethod |
39
|
|
|
def or_none(cls, s: Union[int, str, __qualname__]) -> Optional[__qualname__]: |
|
|
|
|
40
|
|
|
try: |
41
|
|
|
return cls.of(s) |
42
|
|
|
except KeyError: |
43
|
|
|
return None |
44
|
|
|
|
45
|
|
|
@classmethod |
46
|
|
|
def of(cls, s: Union[int, str, __qualname__]) -> __qualname__: |
|
|
|
|
47
|
|
|
""" |
48
|
|
|
Turns a string or int into this type. |
49
|
|
|
Case-insensitive. Replaces `` ``, ``.``, and ``-`` with ``_``. |
50
|
|
|
""" |
51
|
|
|
if isinstance(s, cls): |
52
|
|
|
return s |
53
|
|
|
key = s.strip().replace(" ", "_").replace(".", "_").replace("-", "_").lower() |
54
|
|
|
try: |
55
|
|
|
if isinstance(s, str): |
|
|
|
|
56
|
|
|
return cls[key] |
57
|
|
|
elif isinstance(key, int): |
58
|
|
|
return cls(key) |
59
|
|
|
else: |
60
|
|
|
raise TypeError(f"Lookup type {type(s)} for value {s} not a str or int") |
61
|
|
|
except KeyError: |
62
|
|
|
unk = cls._unmatched_type() |
63
|
|
|
if unk is None: |
64
|
|
|
raise |
65
|
|
|
logger.error(f"Value {key} not found. Using {unk}") |
66
|
|
|
if not isinstance(unk, cls): |
67
|
|
|
raise AssertionError(f"Wrong type {type(unk)} (lookup: {s})") |
68
|
|
|
return unk |
69
|
|
|
|
70
|
|
|
|
71
|
|
|
__all__ = ["TrueFalseUnknown", "CleverEnum"] |
72
|
|
|
|