Rietveld Code Review Tool
Help | Bug tracker | Discussion group | Source code | Sign in
(198090)

Side by Side Diff: Lib/enum.py

Issue 17947: Code, test, and doc review for PEP-0435 Enum
Patch Set: Created 6 years, 7 months ago
Left:
Right:
Use n/p to move between diff chunks; N/P to move between comments. Please Sign in to add in-line comments.
Jump to:
View unified diff | Download patch
« no previous file with comments | « no previous file | Lib/test/test_enum.py » ('j') | Lib/test/test_enum.py » ('J')
Toggle Intra-line Diffs ('i') | Expand Comments ('e') | Collapse Comments ('c') | Show Comments Hide Comments ('s')
OLDNEW
(Empty)
1 """\
2 Provides the Enum class, which can be subclassed to create new, static, enumerat ions.
3 """
4
5 from collections import OrderedDict
6 import sys
7
8 __all__ = ['Enum', 'IntEnum']
9
10
11 class _StealthProperty():
12 """
13 Returns the value in the instance, or the virtual attribute on the class.
14
15 A virtual attribute is one that is looked up by __getattr__, as opposed to
16 one that lives in __class__.__dict__
17
18 """
19
20 def __init__(self, fget=None, fset=None, fdel=None, doc=None):
21 self.fget = fget
22 self.fset = fset
23 self.fdel = fdel
24 self.__doc__ = doc or fget.__doc__
25
26 def __call__(self, func):
27 self.fget = func
28 self.name = func.name
29 if not self.__doc__:
30 self.__doc__ = self.fget.__doc__
31
32 def __get__(self, obj, objtype=None):
33 if obj is None:
34 return getattr(objtype, self.name)
35 if self.fget is None:
36 raise AttributeError("unreadable attribute")
37 return self.fget(obj)
38
39 def __set__(self, obj, value):
40 if self.fset is None:
41 raise AttributeError("can't set attribute")
42 self.fset(obj, value)
43
44 def __delete__(self, obj):
45 if self.fdel is None:
46 raise AttributeError("can't delete attribute")
47 self.fdel(obj)
48
49 def setter(self, func):
50 self.fset = func
51 return self
52
53 def deleter(self, func):
54 self.fdel = func
55 return self
alex 2013/05/10 20:24:06 Real properties don't mutate on setter/deleter/get
stoneleaf 2013/05/12 14:47:01 There's a getter as well?
56
57
58 class _EnumDict(dict):
59 """Keeps track of definition order of the enum items.
60
61 EnumMeta will use the names found in self._enum_names as the
62 enumeration member names."""
63
64 def __init__(self):
65 super().__init__()
66 self._enum_names = []
67
68 def __setitem__(self, key, something):
alex 2013/05/10 20:24:06 I think `value` is a more common name to use than
stoneleaf 2013/05/12 14:47:01 Done.
69 """\
70 Changes anything not dundered or that doesn't have __get__.
71
72 If a descriptor is added with the same name as an enum member, the name
73 is removed from _enum_names (this may leave a hole in the numerical
74 sequence of values).
75
76 If an enum member name is used twice, but the value is not the same, an
77 error is raised.
78
79 """
80 if key[:2] == key[-2:] == '__' or hasattr(something, '__get__'):
81 if key in self._enum_names:
82 # overwriting an enum with a method? then remove the name from
83 # _enum_names or it will become an enum anyway when the class
84 # is created
85 self._enum_names.remove(key)
86 else:
87 if key in self._enum_names and self[key] != something:
88 raise TypeError('Attempted to reuse key: %s' % key)
89 self._enum_names.append(key)
90 super().__setitem__(key, something)
91
92 # dummy value for Enum as EnumMeta explicity checks for it, but of course until
93 # EnumMeta finishes running the first time the Enum class doesn't exist. This
94 # is also why there are checks in EnumMeta like `if Enum is not None`
95 Enum = None
96
97
98 class EnumMeta(type):
99 """\
100 Metaclass for Enum
101 """
102
103 @classmethod
104 def __prepare__(metacls, cls, bases):
105 return _EnumDict()
106
107 def __new__(metacls, cls, bases, classdict):
108 # an Enum class is final once enumeration items have been defined; it
109 # cannot be mixed with other types (int, float, etc.) if it has an
110 # inherited __new__ unless a new __new__ is defined (or the resulting
111 # class will fail).
112 obj_type, first_enum = metacls._get_mixins(bases)
113 __new__, save_new, use_args = metacls._find_new(classdict, obj_type, fir st_enum)
114 # save enum items into separate mapping so they don't get baked into
115 # the new class
116 name_value = {k: classdict[k] for k in classdict._enum_names}
117 for name in classdict._enum_names:
118 del classdict[name]
119 enum_map = OrderedDict()
120 enum_class = type.__new__(metacls, cls, bases, classdict)
121 enum_names= []
122 enum_class._enum_names = enum_names # enum names in definition order
123 enum_class._enum_map = enum_map # name:value map
124 # instantiate them, checking for duplicates as we go
125 # we instantiate first instead of checking for duplicates first in case
126 # a custom __new__ is doing something funky with the values -- such as
127 # auto-numbering ;)
128 for e in classdict._enum_names:
129 value = name_value[e]
130 if not isinstance(value, tuple):
131 args = (value, )
132 else:
133 args = value
134 if obj_type is tuple: # special case for tuple enums
135 args = (args, ) # wrap it one more time
136 if not use_args:
137 enum_item = __new__(enum_class)
138 else:
139 enum_item = __new__(enum_class, *args)
140 enum_item.__init__(*args)
141 enum_item._name = e
142 if not hasattr(enum_item, '_value'):
143 enum_item._value = value
144 # look for any duplicate values, and, if found, use the already
145 # created enum item instead of the new one so `is` will work
146 # (i.e. Color.green is Color.grene)
147 for name, canonical_enum in enum_map.items():
148 if canonical_enum.value == enum_item._value:
149 enum_item = canonical_enum
150 break
151 else:
152 enum_names.append(e)
153 enum_map[e] = enum_item
154 # double check that repr and friends are not the mixin's or various
155 # things break (such as pickle)
156 for name in ('__repr__', '__str__', '__getnewargs__'):
157 class_method = getattr(enum_class, name)
158 obj_method = getattr(obj_type, name, None)
159 enum_method = getattr(first_enum, name, None)
160 if obj_method is not None and obj_method is class_method:
161 setattr(enum_class, name, enum_method)
162 # replace any other __new__ with our own (as long as Enum is not None,
163 # anyway) -- again, this is to support pickle
164 if Enum is not None:
165 # if the user defined their own __new__, save it before it gets
166 # clobbered in case they subclass later
167 if save_new:
168 enum_class.__new_member__ = __new__
169 enum_class.__new__ = Enum.__new__
170 return enum_class
171
172 def __call__(cls, value, names=None, *, module=None, type=None):
173 if names is None: # simple value lookup
174 return cls.__new__(cls, value)
175 # otherwise, we're creating a new Enum type
176 class_name = value # better name for a name than value ;)
177 metacls = cls.__class__
178 bases = (cls, ) if type is None else (type, cls)
179 classdict = metacls.__prepare__(class_name, bases)
180 if isinstance(names, str):
181 names = names.replace(',', ' ').split()
182 if isinstance(names, (tuple, list)) and isinstance(names[0], str):
183 names = [(e, i) for (i, e) in enumerate(names, 1)]
184 # otherwise names better be an iterable of (name, value) or a mapping
185 for item in names:
186 if isinstance(item, str):
187 e, v = item, names[item]
188 else:
189 e, v = item
190 classdict[e] = v
191 enum_class = metacls.__new__(metacls, class_name, bases, classdict)
192 # TODO: replace the frame hack if a blessed way to know the calling
193 # module is ever developed
194 if module is None:
195 try:
196 module = sys._getframe(1).f_globals['__name__']
197 except (AttributeError, ValueError):
198 pass
199 if module is not None:
200 enum_class.__module__ = module
201 return enum_class
202
203 def __contains__(cls, enum_item):
204 return isinstance(enum_item, cls) and enum_item.name in cls._enum_map
205
206 def __dir__(self):
207 return (['__class__', '__doc__', '__members__']
208 + list(self.__members__))
209
210 @property
211 def __members__(cls):
212 return cls._enum_map.copy()
213
214 def __getattr__(cls, name):
215 """Return the enum item matching `name`"""
216 if name[:2] == name[-2:] == '__':
217 raise AttributeError(name)
218 try:
219 return cls._enum_map[name]
220 except KeyError:
221 raise AttributeError(name) from None
222
223 def __getitem__(cls, name):
224 return cls._enum_map[name]
225
226 def __iter__(cls):
227 return (cls._enum_map[name] for name in cls._enum_names)
228
229 def __len__(cls):
230 return len(cls._enum_names)
231
232 def __repr__(cls):
233 return "<enum %r>" % cls.__name__
234
235 @staticmethod
236 def _get_mixins(bases):
237 obj_type = first_enum = None
238 # double check that we are not subclassing a class with existing
239 # enumeration members; while we're at it, see if any other data
240 # type has been mixed in so we can use the correct __new__
241 if bases:
242 for base in bases:
243 if (base is not Enum
244 and issubclass(base, Enum)
245 and base._enum_names):
246 raise TypeError("Cannot extend enumerations")
247 # base is now the last base in bases
248 if not issubclass(base, Enum):
249 raise TypeError("new enumerations must be created as "
250 "`ClassName([mixin_type,] enum_type)`")
251 # get correct mix-in type (either mix-in type of Enum subclass, or
252 # first base if last base is Enum)
253 if not issubclass(bases[0], Enum):
254 obj_type = bases[0] # first data type
255 first_enum = bases[-1] # enum type
256 else:
257 for base in bases[0].__mro__:
258 # most common: (IntEnum, int, Enum, object)
259 # possible: (<Enum 'AutoIntEnum'>, <Enum 'IntEnum'>,
260 # <class 'int'>, <Enum 'Enum'>,
261 # <class 'object'>)
262 if issubclass(base, Enum):
263 if first_enum is None:
264 first_enum = base
265 else:
266 if obj_type is None:
267 obj_type = base
268 else:
269 obj_type = object
270 first_enum = Enum
271 return obj_type, first_enum
272
273 @staticmethod
274 def _find_new(classdict, obj_type, first_enum):
275 # now find the correct __new__, checking to see of one was defined
276 # by the user; also check earlier enum classes in case a __new__ was
277 # saved as __new_member__
278 __new__ = classdict.get('__new__', None)
279 # should __new__ be saved as __new_member__ later?
280 save_new = __new__ is not None
281 if __new__ is None:
282 # check all possibles for __new_member__ before falling back to
283 # __new__
284 for method in ('__new_member__', '__new__'):
285 for possible in (obj_type, first_enum):
286 target = getattr(possible, method, None)
287 if target not in (None,
288 None.__new__,
289 object.__new__,
290 Enum.__new__):
291 __new__ = target
292 break
293 if __new__ is not None:
294 break
295 else:
296 __new__ = object.__new__
297 # if a non-object.__new__ is used then whatever value/tuple was
298 # assigned to the enum member name will be passed to __new__ and to the
299 # new enum member's __init__
300 if __new__ is object.__new__:
301 use_args = False
302 else:
303 use_args = True
304 return __new__, save_new, use_args
305
306
307 class Enum(metaclass=EnumMeta):
308 """valueless, unordered enumeration class"""
309
310 # no actual assignments are made as it is a chicken-and-egg problem
311 # with the metaclass, which checks for the Enum class specifically
312
313 def __new__(cls, value):
314 # all enum instances are actually created during class construction
315 # without calling this method; this method is called by the metaclass'
316 # __call__ (i.e. Color(3) ), and by pickle
317 if type(value) is cls:
318 return value
319 # by-value search for a matching enum member
320 for member in cls._enum_map.values():
321 if member.value == value:
322 return member
323 raise ValueError("%s is not a valid %s" % (value, cls.__name__))
324
325 def __repr__(self):
326 return "<%s.%s: %r>" % (
327 self.__class__.__name__, self._name, self._value)
328
329 def __str__(self):
330 return "%s.%s" % (self.__class__.__name__, self._name)
331
332 def __dir__(self):
333 return (['__class__', '__doc__', 'name', 'value'])
334
335 def __eq__(self, other):
336 if type(other) is self.__class__:
337 return self is other
338 return NotImplemented
339
340 def __getnewargs__(self):
341 return (self._value, )
342
343 def __hash__(self):
344 return hash(self._name)
345
346 # _StealthProperty is used to provide access to the `name` and `value`
347 # properties of enum members while keeping some measure of protection
348 # from modification, while still allowing for an enumeration to have
349 # members named `name` and `value`. This works because enumeration
350 # members are not set directely on the enum class -- __getattr__ is
351 # used to look them up.
352
353 @_StealthProperty
354 def name(self):
355 return self._name
356
357 @_StealthProperty
358 def value(self):
359 return self._value
360
361
362 class IntEnum(int, Enum):
363 """Enum where members are also (and must be) ints"""
364
OLDNEW
« no previous file with comments | « no previous file | Lib/test/test_enum.py » ('j') | Lib/test/test_enum.py » ('J')

RSS Feeds Recent Issues | This issue
This is Rietveld 894c83f36cb7+