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

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 """\
eli.bendersky 2013/05/11 05:23:57 Again, *please* use PEP 257 convention. Even if yo
stoneleaf 2013/05/12 14:47:01 Hadn't gotten that far yet. ;) Is it good now?
2 Provides the Enum class, which can be subclassed to create new, static,
3 enumerations.
4 """
5
6 from collections import OrderedDict
7 import sys
8
9 __all__ = ['Enum', 'IntEnum']
10
11
12 class _StealthProperty():
13 """
eli.bendersky 2013/05/11 05:23:57 PEP 257 everywhere... """First line of multi-line
stoneleaf 2013/05/12 14:47:01 Done.
14 Returns the value in the instance, or the virtual attribute on the class.
15
16 A virtual attribute is one that is looked up by __getattr__, as opposed to
17 one that lives in __class__.__dict__
18
19 """
20
21 def __init__(self, fget=None, fset=None, fdel=None, doc=None):
22 self.fget = fget
23 self.fset = fset
24 self.fdel = fdel
25 self.__doc__ = doc or fget.__doc__
26
27 def __call__(self, func):
28 self.fget = func
29 self.name = func.name
30 if not self.__doc__:
31 self.__doc__ = self.fget.__doc__
32
33 def __get__(self, obj, objtype=None):
34 if obj is None:
35 return getattr(objtype, self.name)
36 if self.fget is None:
37 raise AttributeError("unreadable attribute")
38 return self.fget(obj)
39
40 def __set__(self, obj, value):
41 if self.fset is None:
42 raise AttributeError("can't set attribute")
43 self.fset(obj, value)
44
45 def __delete__(self, obj):
46 if self.fdel is None:
47 raise AttributeError("can't delete attribute")
48 self.fdel(obj)
49
50 def getter(self, func):
51 return self.__class__(self.fget, self.fset, self.fdel, self.__doc__)
52
53 def setter(self, func):
54 return self.__class__(self.fget, self.fset, self.fdel, self.__doc__)
55
56 def deleter(self, func):
57 return self.__class__(self.fget, self.fset, self.fdel, self.__doc__)
alex 2013/05/11 05:08:27 These don't seem to correctly mix the new func in.
stoneleaf 2013/05/12 14:47:01 Sorry, was on auto-pilot.
58
59
60 class _EnumDict(dict):
61 """Keeps track of definition order of the enum items.
62
63 EnumMeta will use the names found in self._enum_names as the
64 enumeration member names."""
65
66 def __init__(self):
67 super().__init__()
68 self._enum_names = []
69
70 def __setitem__(self, key, value):
71 """\
72 Changes anything not dundered or that doesn't have __get__.
73
74 If a descriptor is added with the same name as an enum member, the name
75 is removed from _enum_names (this may leave a hole in the numerical
76 sequence of values).
77
78 If an enum member name is used twice, but the value is not the same, an
79 error is raised.
80
81 """
82 if key[:2] == key[-2:] == '__' or hasattr(value, '__get__'):
83 if key in self._enum_names:
84 # overwriting an enum with a method? then remove the name from
85 # _enum_names or it will become an enum anyway when the class
86 # is created
87 self._enum_names.remove(key)
88 else:
89 if key in self._enum_names:
90 raise TypeError('Attempted to reuse key: %r' % key)
91 self._enum_names.append(key)
92 super().__setitem__(key, value)
93
94 # dummy value for Enum as EnumMeta explicity checks for it, but of course until
95 # EnumMeta finishes running the first time the Enum class doesn't exist. This
96 # is also why there are checks in EnumMeta like `if Enum is not None`
97 Enum = None
98
99
100 class EnumMeta(type):
101 """\
eli.bendersky 2013/05/11 05:23:57 """Metaclass for Enum"""
stoneleaf 2013/05/12 14:47:01 Done.
102 Metaclass for Enum
103 """
104
105 @classmethod
106 def __prepare__(metacls, cls, bases):
107 return _EnumDict()
108
109 def __new__(metacls, cls, bases, classdict):
110 # an Enum class is final once enumeration items have been defined; it
111 # cannot be mixed with other types (int, float, etc.) if it has an
112 # inherited __new__ unless a new __new__ is defined (or the resulting
113 # class will fail).
114 obj_type, first_enum = metacls._get_mixins(bases)
115 __new__, save_new, use_args = metacls._find_new(classdict, obj_type, fir st_enum)
eli.bendersky 2013/05/11 05:23:57 80-column violation
stoneleaf 2013/05/12 14:47:01 Done.
116 # save enum items into separate mapping so they don't get baked into
117 # the new class
118 name_value = {k: classdict[k] for k in classdict._enum_names}
119 for name in classdict._enum_names:
120 del classdict[name]
121 enum_map = OrderedDict()
122 enum_class = type.__new__(metacls, cls, bases, classdict)
123 enum_names= []
124 enum_class._enum_names = enum_names # enum names in definition order
125 enum_class._enum_map = enum_map # name:value map
126 # instantiate them, checking for duplicates as we go
127 # we instantiate first instead of checking for duplicates first in case
128 # a custom __new__ is doing something funky with the values -- such as
129 # auto-numbering ;)
130 for e in classdict._enum_names:
131 value = name_value[e]
132 if not isinstance(value, tuple):
133 args = (value, )
134 else:
135 args = value
136 if obj_type is tuple: # special case for tuple enums
137 args = (args, ) # wrap it one more time
138 if not use_args:
139 enum_item = __new__(enum_class)
140 else:
141 enum_item = __new__(enum_class, *args)
142 enum_item.__init__(*args)
143 enum_item._name = e
144 if not hasattr(enum_item, '_value'):
145 enum_item._value = value
146 # look for any duplicate values, and, if found, use the already
147 # created enum item instead of the new one so `is` will work
148 # (i.e. Color.green is Color.grene)
149 for name, canonical_enum in enum_map.items():
150 if canonical_enum.value == enum_item._value:
151 enum_item = canonical_enum
152 break
153 else:
154 enum_names.append(e)
155 enum_map[e] = enum_item
156 # double check that repr and friends are not the mixin's or various
157 # things break (such as pickle)
158 for name in ('__repr__', '__str__', '__getnewargs__'):
159 class_method = getattr(enum_class, name)
160 obj_method = getattr(obj_type, name, None)
161 enum_method = getattr(first_enum, name, None)
162 if obj_method is not None and obj_method is class_method:
163 setattr(enum_class, name, enum_method)
164 # replace any other __new__ with our own (as long as Enum is not None,
165 # anyway) -- again, this is to support pickle
166 if Enum is not None:
167 # if the user defined their own __new__, save it before it gets
168 # clobbered in case they subclass later
169 if save_new:
170 enum_class.__new_member__ = __new__
171 enum_class.__new__ = Enum.__new__
172 return enum_class
173
174 def __call__(cls, value, names=None, *, module=None, type=None):
eli.bendersky 2013/05/11 05:23:57 Explain the semantics of __call__ in detail in a d
stoneleaf 2013/05/12 14:47:01 Done.
175 if names is None: # simple value lookup
176 return cls.__new__(cls, value)
177 # otherwise, we're creating a new Enum type
178 class_name = value # better name for a name than value ;)
179 metacls = cls.__class__
180 bases = (cls, ) if type is None else (type, cls)
181 classdict = metacls.__prepare__(class_name, bases)
182 if isinstance(names, str):
183 names = names.replace(',', ' ').split()
184 if isinstance(names, (tuple, list)) and isinstance(names[0], str):
185 names = [(e, i) for (i, e) in enumerate(names, 1)]
186 # otherwise names better be an iterable of (name, value) or a mapping
187 for item in names:
188 if isinstance(item, str):
189 e, v = item, names[item]
190 else:
191 e, v = item
192 classdict[e] = v
193 enum_class = metacls.__new__(metacls, class_name, bases, classdict)
194 # TODO: replace the frame hack if a blessed way to know the calling
195 # module is ever developed
196 if module is None:
197 try:
198 module = sys._getframe(1).f_globals['__name__']
199 except (AttributeError, ValueError):
200 pass
201 if module is not None:
202 enum_class.__module__ = module
203 return enum_class
204
205 def __contains__(cls, enum_item):
206 return isinstance(enum_item, cls) and enum_item.name in cls._enum_map
207
208 def __dir__(self):
209 return (['__class__', '__doc__', '__members__'] +
210 list(self.__members__))
211
212 @property
213 def __members__(cls):
214 return cls._enum_map.copy()
eli.bendersky 2013/05/11 05:23:57 I still don't like this copy(). Can you say what i
stoneleaf 2013/05/12 14:47:01 Doc string added. I don't think speed is an issue
eli.bendersky 2013/05/12 15:21:31 I see what you're saying and it makes sense in C++
Nick Coghlan 2013/05/12 15:28:03 This kind of use case is exactly why types.Mapping
eli.bendersky 2013/05/12 15:38:28 Oh cool, I didn't know that existed (*blush*). Loo
stoneleaf 2013/05/13 20:32:37 Thanks, Nick -- that I can do.
215
216 def __getattr__(cls, name):
217 """Return the enum item matching `name`"""
eli.bendersky 2013/05/11 05:23:57 s/item/member/ Use members consistently everywher
stoneleaf 2013/05/12 14:47:01 Done.
218 if name[:2] == name[-2:] == '__':
eli.bendersky 2013/05/11 05:23:57 You have this more than once in the code. Have a h
stoneleaf 2013/05/12 14:47:01 Are you sure? Seems like this would a much bigger
219 raise AttributeError(name)
220 try:
221 return cls._enum_map[name]
222 except KeyError:
223 raise AttributeError(name) from None
alex 2013/05/11 05:08:27 Is there a reason to use __getattr__ over descript
stoneleaf 2013/05/12 14:47:01 Yup, docstring added.
224
225 def __getitem__(cls, name):
226 return cls._enum_map[name]
227
228 def __iter__(cls):
229 return (cls._enum_map[name] for name in cls._enum_names)
230
231 def __len__(cls):
232 return len(cls._enum_names)
233
234 def __repr__(cls):
235 return "<enum %r>" % cls.__name__
236
237 @staticmethod
238 def _get_mixins(bases):
239 if not bases:
240 return object, Enum
241 # double check that we are not subclassing a class with existing
242 # enumeration members; while we're at it, see if any other data
243 # type has been mixed in so we can use the correct __new__
244 obj_type = first_enum = None
245 for base in bases:
246 if (base is not Enum and
247 issubclass(base, Enum) and
248 base._enum_names):
249 raise TypeError("Cannot extend enumerations")
250 # base is now the last base in bases
251 if not issubclass(base, Enum):
252 raise TypeError("new enumerations must be created as "
253 "`ClassName([mixin_type,] enum_type)`")
254 # get correct mix-in type (either mix-in type of Enum subclass, or
255 # first base if last base is Enum)
256 if not issubclass(bases[0], Enum):
257 obj_type = bases[0] # first data type
258 first_enum = bases[-1] # enum type
259 else:
260 for base in bases[0].__mro__:
261 # most common: (IntEnum, int, Enum, object)
262 # possible: (<Enum 'AutoIntEnum'>, <Enum 'IntEnum'>,
263 # <class 'int'>, <Enum 'Enum'>,
264 # <class 'object'>)
265 if issubclass(base, Enum):
266 if first_enum is None:
267 first_enum = base
268 else:
269 if obj_type is None:
270 obj_type = base
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
eli.bendersky 2013/05/11 05:23:57 Here and everywhere - use docstrings to provide a
stoneleaf 2013/05/12 14:47:01 Done.
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+