1
|
|
|
import enum |
|
|
|
|
2
|
|
|
import logging |
3
|
|
|
from typing import AbstractSet |
4
|
|
|
|
5
|
|
|
from pocketutils.core.exceptions import XKeyError |
|
|
|
|
6
|
|
|
|
7
|
|
|
logger = logging.getLogger("pocketutils") |
8
|
|
|
|
9
|
|
|
|
10
|
|
|
class DisjointEnum(enum.Enum): |
11
|
|
|
""" |
12
|
|
|
An enum that does not have combinations. |
13
|
|
|
""" |
14
|
|
|
|
15
|
|
|
@classmethod |
16
|
|
|
def _fix_lookup(cls, s: str) -> str: |
|
|
|
|
17
|
|
|
return s |
18
|
|
|
|
19
|
|
|
@classmethod |
20
|
|
|
def or_none(cls, s: str | __qualname__) -> __qualname__ | None: |
|
|
|
|
21
|
|
|
""" |
22
|
|
|
Returns a choice by name (or returns ``s`` itself). |
23
|
|
|
Returns ``None`` if the choice is not found. |
24
|
|
|
""" |
25
|
|
|
try: |
26
|
|
|
return cls.of(s) |
27
|
|
|
except KeyError: |
28
|
|
|
return None |
29
|
|
|
|
30
|
|
|
@classmethod |
31
|
|
|
def of(cls, s: str | __qualname__) -> __qualname__: |
|
|
|
|
32
|
|
|
""" |
33
|
|
|
Returns a choice by name (or returns ``s`` itself). |
34
|
|
|
""" |
35
|
|
|
if isinstance(s, cls): |
36
|
|
|
return s |
37
|
|
|
return cls[cls._fix_lookup(s)] |
38
|
|
|
|
39
|
|
|
def __new__(cls, *args, **kwargs): |
|
|
|
|
40
|
|
|
value = len(cls.__members__) + 1 |
41
|
|
|
obj = object.__new__(cls) |
42
|
|
|
obj._value_ = value |
43
|
|
|
return obj |
44
|
|
|
|
45
|
|
|
def __repr__(self): |
|
|
|
|
46
|
|
|
return self.name |
47
|
|
|
|
48
|
|
|
def __str__(self): |
|
|
|
|
49
|
|
|
return self.name |
50
|
|
|
|
51
|
|
|
|
52
|
|
|
class FlagEnum(enum.Flag): |
53
|
|
|
""" |
54
|
|
|
A bit flag that behaves as a set, has a null set, and auto-sets values and names. |
55
|
|
|
|
56
|
|
|
Example: |
57
|
|
|
.. code-block:: |
58
|
|
|
|
59
|
|
|
class Flavor(FlagEnum): |
60
|
|
|
NONE = () |
61
|
|
|
BITTER = () |
62
|
|
|
SWEET = () |
63
|
|
|
SOUR = () |
64
|
|
|
UMAMI = () |
65
|
|
|
bittersweet = Flavor.BITTER | Flavor.SWEET |
66
|
|
|
print(bittersweet.value) # 1 + 2 == 3 |
67
|
|
|
print(bittersweet.name) # "bitter|sweet" |
68
|
|
|
|
69
|
|
|
.. important:: |
70
|
|
|
The *first element* must always be the null set ("no flags") |
71
|
|
|
and should be named something like 'none', 'empty', or 'zero' |
72
|
|
|
""" |
73
|
|
|
|
74
|
|
|
@classmethod |
75
|
|
|
def _fix_lookup(cls, s: str) -> str: |
|
|
|
|
76
|
|
|
return s |
77
|
|
|
|
78
|
|
|
def __new__(cls, *args, **kwargs): |
|
|
|
|
79
|
|
|
if len(cls.__members__) == 0: |
80
|
|
|
value = 0 |
81
|
|
|
else: |
82
|
|
|
value = 2 ** (len(cls.__members__) - 1) |
83
|
|
|
obj = object.__new__(cls) |
84
|
|
|
obj._value_ = value |
85
|
|
|
return obj |
86
|
|
|
|
87
|
|
|
@classmethod |
88
|
|
|
def _create_pseudo_member(cls, value): |
89
|
|
|
value = super()._create_pseudo_member(value) |
|
|
|
|
90
|
|
|
members, _ = enum._decompose(cls, value) |
|
|
|
|
91
|
|
|
value._name_ = "|".join([m.name for m in members]) |
|
|
|
|
92
|
|
|
return value |
93
|
|
|
|
94
|
|
|
@classmethod |
95
|
|
|
def or_none(cls, s: str | __qualname__) -> __qualname__ | None: |
|
|
|
|
96
|
|
|
""" |
97
|
|
|
Returns a choice by name (or returns ``s`` itself). |
98
|
|
|
|
99
|
|
|
Returns: |
100
|
|
|
``None`` if the choice is not found. |
101
|
|
|
""" |
102
|
|
|
try: |
103
|
|
|
return cls.of(s) |
104
|
|
|
except KeyError: |
105
|
|
|
return None |
106
|
|
|
|
107
|
|
|
@classmethod |
108
|
|
|
def of(cls, s: str | __qualname__ | AbstractSet[str | __qualname__]) -> __qualname__: |
|
|
|
|
109
|
|
|
""" |
110
|
|
|
Returns a choice by name (or ``s`` itself), or a set of those. |
111
|
|
|
""" |
112
|
|
|
if isinstance(s, cls): |
113
|
|
|
return s |
114
|
|
|
if isinstance(s, str): |
115
|
|
|
return cls[cls._fix_lookup(s)] |
116
|
|
|
z = cls(0) |
|
|
|
|
117
|
|
|
for m in s: |
|
|
|
|
118
|
|
|
z |= cls.of(m) |
|
|
|
|
119
|
|
|
return z |
120
|
|
|
|
121
|
|
|
def __repr__(self): |
|
|
|
|
122
|
|
|
return self.name |
123
|
|
|
|
124
|
|
|
def __str__(self): |
|
|
|
|
125
|
|
|
return self.name |
126
|
|
|
|
127
|
|
|
|
128
|
|
|
class TrueFalseEither(DisjointEnum): |
129
|
|
|
""" |
130
|
|
|
A :class:`pocketutils.core.enums.DisjointEnum` of true, false, or unknown. |
131
|
|
|
""" |
132
|
|
|
|
133
|
|
|
TRUE = () |
134
|
|
|
FALSE = () |
135
|
|
|
EITHER = () |
136
|
|
|
|
137
|
|
|
@classmethod |
138
|
|
|
def _if_not_found(cls, s: str | __qualname__) -> __qualname__: |
|
|
|
|
139
|
|
|
return cls.EITHER |
140
|
|
|
|
141
|
|
|
@classmethod |
142
|
|
|
def _fix_lookup(cls, s: str) -> str: |
143
|
|
|
return s.upper().strip() |
144
|
|
|
|
145
|
|
|
|
146
|
|
|
class MultiTruth(FlagEnum): |
147
|
|
|
""" |
148
|
|
|
A :class:`pocketutils.core.enums.FlagEnum` for true, false, true+false, and neither. |
149
|
|
|
""" |
150
|
|
|
|
151
|
|
|
FALSE = () |
152
|
|
|
TRUE = () |
153
|
|
|
|
154
|
|
|
|
155
|
|
|
class CleverEnum(DisjointEnum): |
156
|
|
|
""" |
157
|
|
|
An enum with a :meth:`of` method that finds values with limited string/value fixing. |
158
|
|
|
Replaces ``" "`` and ``"-"``` with ``_`` and ignores case in :meth:`of`. |
159
|
|
|
May support an "unmatched" type via :meth:`_if_not_found`, |
160
|
|
|
which can return a fallback value when there is no match. |
161
|
|
|
|
162
|
|
|
Example: |
163
|
|
|
class Thing(CleverEnum): |
164
|
|
|
BUILDING = () |
165
|
|
|
OFFICE_SUPPLY = () |
166
|
|
|
POWER_OUTLET = () |
167
|
|
|
|
168
|
|
|
x = Thing.of("power outlet") |
169
|
|
|
|
170
|
|
|
Example: |
171
|
|
|
class Color(CleverEnum): |
172
|
|
|
RED = () |
173
|
|
|
GREEN = () |
174
|
|
|
BLUE = () |
175
|
|
|
OTHER = () |
176
|
|
|
|
177
|
|
|
@classmethod |
178
|
|
|
def _if_not_found(cls, s: str) -> __qualname__: |
179
|
|
|
# raise XValueError(f"No member for value '{s}'", value=s) from None |
180
|
|
|
# ^ |
181
|
|
|
# the default implementation |
182
|
|
|
logger.warning(f"Color {s} unknown; using {cls.OTHER}") |
183
|
|
|
return cls.OTHER |
184
|
|
|
|
185
|
|
|
.. important:: |
186
|
|
|
|
187
|
|
|
If :meth:`_if_not_found` is overridden, it should return a member value. |
188
|
|
|
(In particular, it should never return ``None``.) |
189
|
|
|
|
190
|
|
|
.. important:: |
191
|
|
|
|
192
|
|
|
To use with non-uppercase enum values (e.g. ``Color.red`` instead of ``Color.RED``), |
193
|
|
|
override :meth:`_fix_lookup` with this:: |
194
|
|
|
|
195
|
|
|
@classmethod |
196
|
|
|
def _fix_lookup(cls, s: str) -> str: |
197
|
|
|
return s.strip().replace(" ", "_").replace("-", "_").lower() |
198
|
|
|
# ^ |
199
|
|
|
# changed |
200
|
|
|
""" |
201
|
|
|
|
202
|
|
|
@classmethod |
203
|
|
|
def of(cls, s: str | __qualname__) -> __qualname__: |
204
|
|
|
try: |
205
|
|
|
return super().of(s) |
206
|
|
|
except KeyError: |
207
|
|
|
return cls._if_not_found(s) |
208
|
|
|
|
209
|
|
|
@classmethod |
210
|
|
|
def _if_not_found(cls, s: str | __qualname__) -> __qualname__: |
|
|
|
|
211
|
|
|
raise XKeyError(f"No member for value '{s}'", key=s) from None |
212
|
|
|
|
213
|
|
|
@classmethod |
214
|
|
|
def _fix_lookup(cls, s: str) -> str: |
215
|
|
|
return s.strip().replace(" ", "_").replace("-", "_").upper() |
216
|
|
|
|
217
|
|
|
|
218
|
|
|
__all__ = ["TrueFalseEither", "DisjointEnum", "FlagEnum", "CleverEnum", "MultiTruth"] |
219
|
|
|
|