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

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

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