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

Side by Side Diff: Lib/enum.py

Issue 17947: Code, test, and doc review for PEP-0435 Enum
Patch Set: Created 6 years, 6 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,
3 enumerations.
4 """
5
6 import sys
7 import weakref
Zach Ware 2013/05/17 15:37:11 As pointed out by Donald Stufft on the issue, weak
stoneleaf 2013/05/17 17:40:14 Yeah, I was going to use it, then didn't. Removed
8 from collections import OrderedDict
9 from types import MappingProxyType
10
11 __all__ = ['Enum', 'IntEnum']
12
13
14 class _StealthProperty:
15 """Returns the value in the instance, raises AttributeError on the class.
16
17 if AttributeError is raised, python will automatically use __getattr__ to
18 try and find the attribute.
19
20 """
21
22 def __init__(self, fget=None):
23 self.fget = fget
24
25 def __get__(self, obj, objtype=None):
26 if obj is None:
27 raise AttributeError(self._name)
28 return self.fget(obj)
29
30 def __set__(self, obj, value):
31 raise AttributeError("can't set attribute")
32
33 def __delete__(self, obj):
34 raise AttributeError("can't delete attribute")
35
36
37 def dunder(name):
Zach Ware 2013/05/17 15:37:11 Might be better named "is_dunder" since you're ask
stoneleaf 2013/05/17 17:40:14 Done.
38 """Returns True if a dunder name, False otherwise."""
39
40 return name[:2] == name[-2:] == '__'
41
42
43 def break_noisily_on_pickle(self):
44 """Raises an exception when trying to pickle the class.
45
46 If module name was not able to be determined using the functional syntax,
47 this function will be stored on that class as __reduce__. It's only
48 purpose is to move the pickle error from /un/pickling to pickling.
49
50 """
51
52 raise TypeError('%r cannot be pickled' % self)
Zach Ware 2013/05/17 15:37:11 These two functions should be made private unless
stoneleaf 2013/05/17 17:40:14 Done.
53
54
55 class _EnumDict(dict):
56 """Keeps track of definition order of the enum items.
57
58 EnumMeta will use the names found in self._enum_names as the
59 enumeration member names.
60
61 """
62
63 def __init__(self):
64 super().__init__()
65 self._enum_names = []
66
67 def __setitem__(self, key, value):
68 """Changes anything not dundered or that doesn't have __get__.
69
70 If a descriptor is added with the same name as an enum member, the name
71 is removed from _enum_names (this may leave a hole in the numerical
72 sequence of values).
73
74 If an enum member name is used twice, an error is raised; duplicate
75 values are not checked for.
76
77 """
78
79 if dunder(key) or hasattr(value, '__get__'):
80 if key in self._enum_names:
81 # overwriting an enum with a method? then remove the name from
82 # _enum_names or it will become an enum anyway when the class
83 # is created
84 self._enum_names.remove(key)
85 else:
86 if key in self._enum_names:
87 raise TypeError('Attempted to reuse key: %r' % key)
88 self._enum_names.append(key)
89 super().__setitem__(key, value)
90
91 # dummy value for Enum as EnumMeta explicity checks for it, but of course until
92 # EnumMeta finishes running the first time the Enum class doesn't exist. This
93 # is also why there are checks in EnumMeta like `if Enum is not None`
94 Enum = None
95
96
97 class EnumMeta(type):
98 """Metaclass for Enum"""
99
100 @classmethod
101 def __prepare__(metacls, cls, bases):
102 return _EnumDict()
103
104 def __new__(metacls, cls, bases, classdict):
105 # an Enum class is final once enumeration items have been defined; it
106 # cannot be mixed with other types (int, float, etc.) if it has an
107 # inherited __new__ unless a new __new__ is defined (or the resulting
108 # class will fail).
109 obj_type, first_enum = metacls._get_mixins(bases)
110 __new__, save_new, use_args = (
111 metacls._find_new(classdict, obj_type, first_enum)
112 )
113
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
120 # check for illegal enum names (any others?)
121 if set(name_value) & {'mro', '_create', '_get_mixins', '_find_new'}:
122 raise ValueError("The names 'mro', '_create', '_get_mixins', and "
123 "'_find_new' cannot be used for members")
Zach Ware 2013/05/17 15:37:11 I wonder if it would be better to explicitly disal
stoneleaf 2013/05/17 17:40:14 I'd rather not -- seems like it would be easy to m
Zach Ware 2013/05/17 18:01:39 That would be fine with me, I'm not aware of any e
124
125 # add _name attributes to any _StealthProperty's for nicer exceptions
126 for name, value in classdict.items():
127 if isinstance(value, _StealthProperty):
128 value._name = name
129
130 # create our new Enum type
131 enum_class = super().__new__(metacls, cls, bases, classdict)
132 enum_names = []
133 enum_map = OrderedDict()
134 enum_value_map = {} # hashable value:name map
135 enum_class._enum_names = enum_names # names in definition order
136 enum_class._enum_map = enum_map # name:value map
137 enum_class._enum_value_map = enum_value_map # value:name map
138
139 # check for a __getnewargs__, and if not present sabotage
140 # pickling, since it won't work anyway
141 if (
142 obj_type is not object and
143 obj_type.__dict__.get('__getnewargs__') is None
144 ):
145 enum_class.__reduce__ = break_noisily_on_pickle
146 enum_class.__module__ = 'uh uh'
alex 2013/05/17 15:59:33 Can we use a slightly more formal value here?
stoneleaf 2013/05/17 17:40:14 If you insist. ;) Went with 'not picklable'.
147
148 # instantiate them, checking for duplicates as we go
149 # we instantiate first instead of checking for duplicates first in case
150 # a custom __new__ is doing something funky with the values -- such as
151 # auto-numbering ;)
152 for e in classdict._enum_names:
153 value = name_value[e]
154 if not isinstance(value, tuple):
155 args = (value, )
156 else:
157 args = value
158 if obj_type is tuple: # special case for tuple enums
159 args = (args, ) # wrap it one more time
160 if not use_args:
161 enum_item = __new__(enum_class)
162 enum_item._value = value
163 else:
164 enum_item = __new__(enum_class, *args)
165 if not hasattr(enum_item, '_value'):
166 enum_item._value = obj_type(*args)
167 enum_item._obj_type = obj_type
168 enum_item._name = e
169 enum_item.__init__(*args)
170 # look for any duplicate values, and, if found, use the already
171 # created enum item instead of the new one so `is` will work
172 # (i.e. Color.green is Color.grene)
173 for name, canonical_enum in enum_map.items():
174 if canonical_enum.value == enum_item._value:
175 enum_item = canonical_enum
176 break
177 else:
178 enum_names.append(e)
179 enum_map[e] = enum_item
180 try:
181 enum_value_map[value] = enum_item
182 except TypeError:
183 pass
184
185 # double check that repr and friends are not the mixin's or various
186 # things break (such as pickle)
187 for name in ('__repr__', '__str__', '__getnewargs__'):
188 class_method = getattr(enum_class, name)
189 obj_method = getattr(obj_type, name, None)
190 enum_method = getattr(first_enum, name, None)
191 if obj_method is not None and obj_method is class_method:
192 setattr(enum_class, name, enum_method)
193
194 # replace any other __new__ with our own (as long as Enum is not None,
195 # anyway) -- again, this is to support pickle
196 if Enum is not None:
197 # if the user defined their own __new__, save it before it gets
198 # clobbered in case they subclass later
199 if save_new:
200 enum_class.__new_member__ = __new__
201 enum_class.__new__ = Enum.__new__
202 return enum_class
203
204 def __call__(cls, value, names=None, *, module=None, type=None):
205 """Either returns an existing member, or creates a new enum class.
206
207 This method is used both when an enum class is given a value to match
208 to an enumeration member (i.e. Color(3)) and for the functional API
209 (i.e. Color = Enum('Color', names='red green blue')).
210
211 When used as the functional API module, if set, will be stored in the
212 new class' __module__ attribute; type, if set, will be mixed in as the
213 first base class.
214
215 Note: if module is not set this routine will attempt to discover the
216 calling module by walking the frame stack; if this is unsuccessful
217 the resulting class will not be pickleable.
218
219 """
220 if names is None: # simple value lookup
221 return cls.__new__(cls, value)
222 # otherwise, we're creating a new Enum type
223 return cls._create(value, names, module=module, type=type)
224
225 def __contains__(cls, enum_item):
226 return isinstance(enum_item, cls) and enum_item.name in cls._enum_map
227
228 def __dir__(self):
229 return ['__class__', '__doc__', '__members__'] + self._enum_names
230
231 @property
232 def __members__(cls):
233 """Returns a MappingProxyType of the internal _enum_map structure."""
234
235 return MappingProxyType(cls._enum_map)
236
237 def __getattr__(cls, name):
238 """Return the enum member matching `name`
239
240 We use __getattr__ instead of descriptors or inserting into the enum
241 class' __dict__ in order to support `name` and `value` being both
242 properties for enum members (which live in the class' __dict__) and
243 enum members themselves.
244
245 """
246
247 if dunder(name):
248 raise AttributeError(name)
249 try:
250 return cls._enum_map[name]
251 except KeyError:
252 raise AttributeError(name) from None
253
254 def __getitem__(cls, name):
255 return cls._enum_map[name]
256
257 def __iter__(cls):
258 return (cls._enum_map[name] for name in cls._enum_names)
259
260 def __len__(cls):
261 return len(cls._enum_names)
262
263 def __repr__(cls):
264 return "<enum %r>" % cls.__name__
265
266 def _create(cls, class_name, names=None, *, module=None, type=None):
267 """Convenience method to create a new Enum class.
268
269 Called by __new__, with the same arguments, to provide the
270 implementation. Easier to subclass this way.
271
272 """
273 metacls = cls.__class__
274 bases = (cls, ) if type is None else (type, cls)
275 classdict = metacls.__prepare__(class_name, bases)
276
277 # special processing needed for names?
278 if isinstance(names, str):
279 names = names.replace(',', ' ').split()
280 if isinstance(names, (tuple, list)) and isinstance(names[0], str):
281 names = [(e, i) for (i, e) in enumerate(names, 1)]
282
283 # otherwise names better be an iterable of (name, value) or a mapping
284 for item in names:
285 if isinstance(item, str):
286 e, v = item, names[item]
287 else:
288 e, v = item
289 classdict[e] = v
290 enum_class = metacls.__new__(metacls, class_name, bases, classdict)
291
292 # TODO: replace the frame hack if a blessed way to know the calling
293 # module is ever developed
294 if module is None:
295 try:
296 module = sys._getframe(2).f_globals['__name__']
297 except (AttributeError, ValueError) as exc:
298 pass
299 if module is None:
300 enum_class.__module__ = 'uh uh'
301 enum_class.__reduce__ = break_noisily_on_pickle
302 else:
303 enum_class.__module__ = module
304
305 return enum_class
306
307 @staticmethod
308 def _get_mixins(bases):
309 """Returns the type for creating enum members, and the first inherited
310 enum class.
311
312 bases: the tuple of bases that was given to __new__
313
314 """
315 if not bases:
316 return object, Enum
317
318 # double check that we are not subclassing a class with existing
319 # enumeration members; while we're at it, see if any other data
320 # type has been mixed in so we can use the correct __new__
321 obj_type = first_enum = None
322 for base in bases:
323 if (base is not Enum and
324 issubclass(base, Enum) and
325 base._enum_names):
326 raise TypeError("Cannot extend enumerations")
327 # base is now the last base in bases
328 if not issubclass(base, Enum):
329 raise TypeError("new enumerations must be created as "
330 "`ClassName([mixin_type,] enum_type)`")
331
332 # get correct mix-in type (either mix-in type of Enum subclass, or
333 # first base if last base is Enum)
334 if not issubclass(bases[0], Enum):
335 obj_type = bases[0] # first data type
336 first_enum = bases[-1] # enum type
337 else:
338 for base in bases[0].__mro__:
339 # most common: (IntEnum, int, Enum, object)
340 # possible: (<Enum 'AutoIntEnum'>, <Enum 'IntEnum'>,
341 # <class 'int'>, <Enum 'Enum'>,
342 # <class 'object'>)
343 if issubclass(base, Enum):
344 if first_enum is None:
345 first_enum = base
346 else:
347 if obj_type is None:
348 obj_type = base
349
350 return obj_type, first_enum
351
352 @staticmethod
353 def _find_new(classdict, obj_type, first_enum):
354 """Returns the __new__ to be used for creating the enum members.
355
356 classdict: the class dictionary given to __new__
357 obj_type: the data type whose __new__ will be used by default
358 first_enum: enumeration to check for an overriding __new__
359
360 """
361 # now find the correct __new__, checking to see of one was defined
362 # by the user; also check earlier enum classes in case a __new__ was
363 # saved as __new_member__
364 __new__ = classdict.get('__new__', None)
365
366 # should __new__ be saved as __new_member__ later?
367 save_new = __new__ is not None
368
369 if __new__ is None:
370 # check all possibles for __new_member__ before falling back to
371 # __new__
372 for method in ('__new_member__', '__new__'):
373 for possible in (obj_type, first_enum):
374 target = getattr(possible, method, None)
375 if target not in {
376 None,
377 None.__new__,
378 object.__new__,
379 Enum.__new__,
380 }:
381 __new__ = target
382 break
383 if __new__ is not None:
384 break
385 else:
386 __new__ = object.__new__
387
388 # if a non-object.__new__ is used then whatever value/tuple was
389 # assigned to the enum member name will be passed to __new__ and to the
390 # new enum member's __init__
391 if __new__ is object.__new__:
392 use_args = False
393 else:
394 use_args = True
395
396 return __new__, save_new, use_args
397
398
399 class Enum(metaclass=EnumMeta):
400 """Valueless, unordered enumeration class"""
401
402 # no actual assignments are made as it is a chicken-and-egg problem
403 # with the metaclass, which checks for the Enum class specifically
404
405 def __new__(cls, value):
406 # all enum instances are actually created during class construction
407 # without calling this method; this method is called by the metaclass'
408 # __call__ (i.e. Color(3) ), and by pickle
409 if type(value) is cls:
410 return value
411 # by-value search for a matching enum member
412 # see if it's in the reverse mapping (for hashable values)
413 if value in cls._enum_value_map:
414 return cls._enum_value_map[value]
415 # not there, now do long search -- O(n) behavior
416 for member in cls._enum_map.values():
417 if member.value == value:
418 return member
419 raise ValueError("%s is not a valid %s" % (value, cls.__name__))
420
421 def __repr__(self):
422 return "<%s.%s: %r>" % (
423 self.__class__.__name__, self._name, self._value)
424
425 def __str__(self):
426 return "%s.%s" % (self.__class__.__name__, self._name)
427
428 def __dir__(self):
429 return (['__class__', '__doc__', 'name', 'value'])
430
431 def __eq__(self, other):
432 if type(other) is self.__class__:
433 return self is other
434 return NotImplemented
435
436 def __getnewargs__(self):
437 return (self._value, )
438
439 def __hash__(self):
440 return hash(self._name)
441
442 # _StealthProperty is used to provide access to the `name` and `value`
443 # properties of enum members while keeping some measure of protection
444 # from modification, while still allowing for an enumeration to have
445 # members named `name` and `value`. This works because enumeration
446 # members are not set directely on the enum class -- __getattr__ is
447 # used to look them up.
448
449 @_StealthProperty
450 def name(self):
451 return self._name
452
453 @_StealthProperty
454 def value(self):
455 return self._value
456
457
458 class IntEnum(int, Enum):
459 """Enum where members are also (and must be) ints"""
460
461
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+