blob: d0b4a1c089247e3a97c9a40cb2a6b06438d0cfa3 [file] [log] [blame]
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001import enum
2import unittest
3from collections import OrderedDict
4from pickle import dumps, loads, PicklingError
Ethan Furmanf24bb352013-07-18 17:05:39 -07005from enum import Enum, IntEnum, unique
Ethan Furman6b3d64a2013-06-14 16:55:46 -07006
7# for pickle tests
8try:
9 class Stooges(Enum):
10 LARRY = 1
11 CURLY = 2
12 MOE = 3
13except Exception as exc:
14 Stooges = exc
15
16try:
17 class IntStooges(int, Enum):
18 LARRY = 1
19 CURLY = 2
20 MOE = 3
21except Exception as exc:
22 IntStooges = exc
23
24try:
25 class FloatStooges(float, Enum):
26 LARRY = 1.39
27 CURLY = 2.72
28 MOE = 3.142596
29except Exception as exc:
30 FloatStooges = exc
31
32# for pickle test and subclass tests
33try:
34 class StrEnum(str, Enum):
35 'accepts only string values'
36 class Name(StrEnum):
37 BDFL = 'Guido van Rossum'
38 FLUFL = 'Barry Warsaw'
39except Exception as exc:
40 Name = exc
41
42try:
43 Question = Enum('Question', 'who what when where why', module=__name__)
44except Exception as exc:
45 Question = exc
46
47try:
48 Answer = Enum('Answer', 'him this then there because')
49except Exception as exc:
50 Answer = exc
51
52# for doctests
53try:
54 class Fruit(Enum):
55 tomato = 1
56 banana = 2
57 cherry = 3
58except Exception:
59 pass
60
61class TestEnum(unittest.TestCase):
62 def setUp(self):
63 class Season(Enum):
64 SPRING = 1
65 SUMMER = 2
66 AUTUMN = 3
67 WINTER = 4
68 self.Season = Season
69
70 def test_enum_in_enum_out(self):
71 Season = self.Season
72 self.assertIs(Season(Season.WINTER), Season.WINTER)
73
74 def test_enum_value(self):
75 Season = self.Season
76 self.assertEqual(Season.SPRING.value, 1)
77
78 def test_intenum_value(self):
79 self.assertEqual(IntStooges.CURLY.value, 2)
80
81 def test_dir_on_class(self):
82 Season = self.Season
83 self.assertEqual(
84 set(dir(Season)),
85 set(['__class__', '__doc__', '__members__',
86 'SPRING', 'SUMMER', 'AUTUMN', 'WINTER']),
87 )
88
89 def test_dir_on_item(self):
90 Season = self.Season
91 self.assertEqual(
92 set(dir(Season.WINTER)),
93 set(['__class__', '__doc__', 'name', 'value']),
94 )
95
96 def test_enum(self):
97 Season = self.Season
98 lst = list(Season)
99 self.assertEqual(len(lst), len(Season))
100 self.assertEqual(len(Season), 4, Season)
101 self.assertEqual(
102 [Season.SPRING, Season.SUMMER, Season.AUTUMN, Season.WINTER], lst)
103
104 for i, season in enumerate('SPRING SUMMER AUTUMN WINTER'.split(), 1):
105 e = Season(i)
106 self.assertEqual(e, getattr(Season, season))
107 self.assertEqual(e.value, i)
108 self.assertNotEqual(e, i)
109 self.assertEqual(e.name, season)
110 self.assertIn(e, Season)
111 self.assertIs(type(e), Season)
112 self.assertIsInstance(e, Season)
113 self.assertEqual(str(e), 'Season.' + season)
114 self.assertEqual(
115 repr(e),
116 '<Season.{0}: {1}>'.format(season, i),
117 )
118
119 def test_value_name(self):
120 Season = self.Season
121 self.assertEqual(Season.SPRING.name, 'SPRING')
122 self.assertEqual(Season.SPRING.value, 1)
123 with self.assertRaises(AttributeError):
124 Season.SPRING.name = 'invierno'
125 with self.assertRaises(AttributeError):
126 Season.SPRING.value = 2
127
128 def test_invalid_names(self):
129 with self.assertRaises(ValueError):
130 class Wrong(Enum):
131 mro = 9
132 with self.assertRaises(ValueError):
133 class Wrong(Enum):
134 _create_= 11
135 with self.assertRaises(ValueError):
136 class Wrong(Enum):
137 _get_mixins_ = 9
138 with self.assertRaises(ValueError):
139 class Wrong(Enum):
140 _find_new_ = 1
141 with self.assertRaises(ValueError):
142 class Wrong(Enum):
143 _any_name_ = 9
144
145 def test_contains(self):
146 Season = self.Season
147 self.assertIn(Season.AUTUMN, Season)
148 self.assertNotIn(3, Season)
149
150 val = Season(3)
151 self.assertIn(val, Season)
152
153 class OtherEnum(Enum):
154 one = 1; two = 2
155 self.assertNotIn(OtherEnum.two, Season)
156
157 def test_comparisons(self):
158 Season = self.Season
159 with self.assertRaises(TypeError):
160 Season.SPRING < Season.WINTER
161 with self.assertRaises(TypeError):
162 Season.SPRING > 4
163
164 self.assertNotEqual(Season.SPRING, 1)
165
166 class Part(Enum):
167 SPRING = 1
168 CLIP = 2
169 BARREL = 3
170
171 self.assertNotEqual(Season.SPRING, Part.SPRING)
172 with self.assertRaises(TypeError):
173 Season.SPRING < Part.CLIP
174
175 def test_enum_duplicates(self):
176 class Season(Enum):
177 SPRING = 1
178 SUMMER = 2
179 AUTUMN = FALL = 3
180 WINTER = 4
181 ANOTHER_SPRING = 1
182 lst = list(Season)
183 self.assertEqual(
184 lst,
185 [Season.SPRING, Season.SUMMER,
186 Season.AUTUMN, Season.WINTER,
187 ])
188 self.assertIs(Season.FALL, Season.AUTUMN)
189 self.assertEqual(Season.FALL.value, 3)
190 self.assertEqual(Season.AUTUMN.value, 3)
191 self.assertIs(Season(3), Season.AUTUMN)
192 self.assertIs(Season(1), Season.SPRING)
193 self.assertEqual(Season.FALL.name, 'AUTUMN')
194 self.assertEqual(
195 [k for k,v in Season.__members__.items() if v.name != k],
196 ['FALL', 'ANOTHER_SPRING'],
197 )
198
199 def test_enum_with_value_name(self):
200 class Huh(Enum):
201 name = 1
202 value = 2
203 self.assertEqual(
204 list(Huh),
205 [Huh.name, Huh.value],
206 )
207 self.assertIs(type(Huh.name), Huh)
208 self.assertEqual(Huh.name.name, 'name')
209 self.assertEqual(Huh.name.value, 1)
210 def test_hash(self):
211 Season = self.Season
212 dates = {}
213 dates[Season.WINTER] = '1225'
214 dates[Season.SPRING] = '0315'
215 dates[Season.SUMMER] = '0704'
216 dates[Season.AUTUMN] = '1031'
217 self.assertEqual(dates[Season.AUTUMN], '1031')
218
219 def test_intenum_from_scratch(self):
220 class phy(int, Enum):
221 pi = 3
222 tau = 2 * pi
223 self.assertTrue(phy.pi < phy.tau)
224
225 def test_intenum_inherited(self):
226 class IntEnum(int, Enum):
227 pass
228 class phy(IntEnum):
229 pi = 3
230 tau = 2 * pi
231 self.assertTrue(phy.pi < phy.tau)
232
233 def test_floatenum_from_scratch(self):
234 class phy(float, Enum):
235 pi = 3.141596
236 tau = 2 * pi
237 self.assertTrue(phy.pi < phy.tau)
238
239 def test_floatenum_inherited(self):
240 class FloatEnum(float, Enum):
241 pass
242 class phy(FloatEnum):
243 pi = 3.141596
244 tau = 2 * pi
245 self.assertTrue(phy.pi < phy.tau)
246
247 def test_strenum_from_scratch(self):
248 class phy(str, Enum):
249 pi = 'Pi'
250 tau = 'Tau'
251 self.assertTrue(phy.pi < phy.tau)
252
253 def test_strenum_inherited(self):
254 class StrEnum(str, Enum):
255 pass
256 class phy(StrEnum):
257 pi = 'Pi'
258 tau = 'Tau'
259 self.assertTrue(phy.pi < phy.tau)
260
261
262 def test_intenum(self):
263 class WeekDay(IntEnum):
264 SUNDAY = 1
265 MONDAY = 2
266 TUESDAY = 3
267 WEDNESDAY = 4
268 THURSDAY = 5
269 FRIDAY = 6
270 SATURDAY = 7
271
272 self.assertEqual(['a', 'b', 'c'][WeekDay.MONDAY], 'c')
273 self.assertEqual([i for i in range(WeekDay.TUESDAY)], [0, 1, 2])
274
275 lst = list(WeekDay)
276 self.assertEqual(len(lst), len(WeekDay))
277 self.assertEqual(len(WeekDay), 7)
278 target = 'SUNDAY MONDAY TUESDAY WEDNESDAY THURSDAY FRIDAY SATURDAY'
279 target = target.split()
280 for i, weekday in enumerate(target, 1):
281 e = WeekDay(i)
282 self.assertEqual(e, i)
283 self.assertEqual(int(e), i)
284 self.assertEqual(e.name, weekday)
285 self.assertIn(e, WeekDay)
286 self.assertEqual(lst.index(e)+1, i)
287 self.assertTrue(0 < e < 8)
288 self.assertIs(type(e), WeekDay)
289 self.assertIsInstance(e, int)
290 self.assertIsInstance(e, Enum)
291
292 def test_intenum_duplicates(self):
293 class WeekDay(IntEnum):
294 SUNDAY = 1
295 MONDAY = 2
296 TUESDAY = TEUSDAY = 3
297 WEDNESDAY = 4
298 THURSDAY = 5
299 FRIDAY = 6
300 SATURDAY = 7
301 self.assertIs(WeekDay.TEUSDAY, WeekDay.TUESDAY)
302 self.assertEqual(WeekDay(3).name, 'TUESDAY')
303 self.assertEqual([k for k,v in WeekDay.__members__.items()
304 if v.name != k], ['TEUSDAY', ])
305
306 def test_pickle_enum(self):
307 if isinstance(Stooges, Exception):
308 raise Stooges
309 self.assertIs(Stooges.CURLY, loads(dumps(Stooges.CURLY)))
310 self.assertIs(Stooges, loads(dumps(Stooges)))
311
312 def test_pickle_int(self):
313 if isinstance(IntStooges, Exception):
314 raise IntStooges
315 self.assertIs(IntStooges.CURLY, loads(dumps(IntStooges.CURLY)))
316 self.assertIs(IntStooges, loads(dumps(IntStooges)))
317
318 def test_pickle_float(self):
319 if isinstance(FloatStooges, Exception):
320 raise FloatStooges
321 self.assertIs(FloatStooges.CURLY, loads(dumps(FloatStooges.CURLY)))
322 self.assertIs(FloatStooges, loads(dumps(FloatStooges)))
323
324 def test_pickle_enum_function(self):
325 if isinstance(Answer, Exception):
326 raise Answer
327 self.assertIs(Answer.him, loads(dumps(Answer.him)))
328 self.assertIs(Answer, loads(dumps(Answer)))
329
330 def test_pickle_enum_function_with_module(self):
331 if isinstance(Question, Exception):
332 raise Question
333 self.assertIs(Question.who, loads(dumps(Question.who)))
334 self.assertIs(Question, loads(dumps(Question)))
335
336 def test_exploding_pickle(self):
337 BadPickle = Enum('BadPickle', 'dill sweet bread-n-butter')
338 enum._make_class_unpicklable(BadPickle)
339 globals()['BadPickle'] = BadPickle
340 with self.assertRaises(TypeError):
341 dumps(BadPickle.dill)
342 with self.assertRaises(PicklingError):
343 dumps(BadPickle)
344
345 def test_string_enum(self):
346 class SkillLevel(str, Enum):
347 master = 'what is the sound of one hand clapping?'
348 journeyman = 'why did the chicken cross the road?'
349 apprentice = 'knock, knock!'
350 self.assertEqual(SkillLevel.apprentice, 'knock, knock!')
351
352 def test_getattr_getitem(self):
353 class Period(Enum):
354 morning = 1
355 noon = 2
356 evening = 3
357 night = 4
358 self.assertIs(Period(2), Period.noon)
359 self.assertIs(getattr(Period, 'night'), Period.night)
360 self.assertIs(Period['morning'], Period.morning)
361
362 def test_getattr_dunder(self):
363 Season = self.Season
364 self.assertTrue(getattr(Season, '__eq__'))
365
366 def test_iteration_order(self):
367 class Season(Enum):
368 SUMMER = 2
369 WINTER = 4
370 AUTUMN = 3
371 SPRING = 1
372 self.assertEqual(
373 list(Season),
374 [Season.SUMMER, Season.WINTER, Season.AUTUMN, Season.SPRING],
375 )
376
377 def test_programatic_function_string(self):
378 SummerMonth = Enum('SummerMonth', 'june july august')
379 lst = list(SummerMonth)
380 self.assertEqual(len(lst), len(SummerMonth))
381 self.assertEqual(len(SummerMonth), 3, SummerMonth)
382 self.assertEqual(
383 [SummerMonth.june, SummerMonth.july, SummerMonth.august],
384 lst,
385 )
386 for i, month in enumerate('june july august'.split(), 1):
387 e = SummerMonth(i)
388 self.assertEqual(int(e.value), i)
389 self.assertNotEqual(e, i)
390 self.assertEqual(e.name, month)
391 self.assertIn(e, SummerMonth)
392 self.assertIs(type(e), SummerMonth)
393
394 def test_programatic_function_string_list(self):
395 SummerMonth = Enum('SummerMonth', ['june', 'july', 'august'])
396 lst = list(SummerMonth)
397 self.assertEqual(len(lst), len(SummerMonth))
398 self.assertEqual(len(SummerMonth), 3, SummerMonth)
399 self.assertEqual(
400 [SummerMonth.june, SummerMonth.july, SummerMonth.august],
401 lst,
402 )
403 for i, month in enumerate('june july august'.split(), 1):
404 e = SummerMonth(i)
405 self.assertEqual(int(e.value), i)
406 self.assertNotEqual(e, i)
407 self.assertEqual(e.name, month)
408 self.assertIn(e, SummerMonth)
409 self.assertIs(type(e), SummerMonth)
410
411 def test_programatic_function_iterable(self):
412 SummerMonth = Enum(
413 'SummerMonth',
414 (('june', 1), ('july', 2), ('august', 3))
415 )
416 lst = list(SummerMonth)
417 self.assertEqual(len(lst), len(SummerMonth))
418 self.assertEqual(len(SummerMonth), 3, SummerMonth)
419 self.assertEqual(
420 [SummerMonth.june, SummerMonth.july, SummerMonth.august],
421 lst,
422 )
423 for i, month in enumerate('june july august'.split(), 1):
424 e = SummerMonth(i)
425 self.assertEqual(int(e.value), i)
426 self.assertNotEqual(e, i)
427 self.assertEqual(e.name, month)
428 self.assertIn(e, SummerMonth)
429 self.assertIs(type(e), SummerMonth)
430
431 def test_programatic_function_from_dict(self):
432 SummerMonth = Enum(
433 'SummerMonth',
434 OrderedDict((('june', 1), ('july', 2), ('august', 3)))
435 )
436 lst = list(SummerMonth)
437 self.assertEqual(len(lst), len(SummerMonth))
438 self.assertEqual(len(SummerMonth), 3, SummerMonth)
439 self.assertEqual(
440 [SummerMonth.june, SummerMonth.july, SummerMonth.august],
441 lst,
442 )
443 for i, month in enumerate('june july august'.split(), 1):
444 e = SummerMonth(i)
445 self.assertEqual(int(e.value), i)
446 self.assertNotEqual(e, i)
447 self.assertEqual(e.name, month)
448 self.assertIn(e, SummerMonth)
449 self.assertIs(type(e), SummerMonth)
450
451 def test_programatic_function_type(self):
452 SummerMonth = Enum('SummerMonth', 'june july august', type=int)
453 lst = list(SummerMonth)
454 self.assertEqual(len(lst), len(SummerMonth))
455 self.assertEqual(len(SummerMonth), 3, SummerMonth)
456 self.assertEqual(
457 [SummerMonth.june, SummerMonth.july, SummerMonth.august],
458 lst,
459 )
460 for i, month in enumerate('june july august'.split(), 1):
461 e = SummerMonth(i)
462 self.assertEqual(e, i)
463 self.assertEqual(e.name, month)
464 self.assertIn(e, SummerMonth)
465 self.assertIs(type(e), SummerMonth)
466
467 def test_programatic_function_type_from_subclass(self):
468 SummerMonth = IntEnum('SummerMonth', 'june july august')
469 lst = list(SummerMonth)
470 self.assertEqual(len(lst), len(SummerMonth))
471 self.assertEqual(len(SummerMonth), 3, SummerMonth)
472 self.assertEqual(
473 [SummerMonth.june, SummerMonth.july, SummerMonth.august],
474 lst,
475 )
476 for i, month in enumerate('june july august'.split(), 1):
477 e = SummerMonth(i)
478 self.assertEqual(e, i)
479 self.assertEqual(e.name, month)
480 self.assertIn(e, SummerMonth)
481 self.assertIs(type(e), SummerMonth)
482
483 def test_subclassing(self):
484 if isinstance(Name, Exception):
485 raise Name
486 self.assertEqual(Name.BDFL, 'Guido van Rossum')
487 self.assertTrue(Name.BDFL, Name('Guido van Rossum'))
488 self.assertIs(Name.BDFL, getattr(Name, 'BDFL'))
489 self.assertIs(Name.BDFL, loads(dumps(Name.BDFL)))
490
491 def test_extending(self):
492 class Color(Enum):
493 red = 1
494 green = 2
495 blue = 3
496 with self.assertRaises(TypeError):
497 class MoreColor(Color):
498 cyan = 4
499 magenta = 5
500 yellow = 6
501
502 def test_exclude_methods(self):
503 class whatever(Enum):
504 this = 'that'
505 these = 'those'
506 def really(self):
507 return 'no, not %s' % self.value
508 self.assertIsNot(type(whatever.really), whatever)
509 self.assertEqual(whatever.this.really(), 'no, not that')
510
511 def test_overwrite_enums(self):
512 class Why(Enum):
513 question = 1
514 answer = 2
515 propisition = 3
516 def question(self):
517 print(42)
518 self.assertIsNot(type(Why.question), Why)
Ethan Furman520ad572013-07-19 19:47:21 -0700519 self.assertNotIn(Why.question, Why._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700520 self.assertNotIn(Why.question, Why)
521
522 def test_wrong_inheritance_order(self):
523 with self.assertRaises(TypeError):
524 class Wrong(Enum, str):
525 NotHere = 'error before this point'
526
527 def test_intenum_transitivity(self):
528 class number(IntEnum):
529 one = 1
530 two = 2
531 three = 3
532 class numero(IntEnum):
533 uno = 1
534 dos = 2
535 tres = 3
536 self.assertEqual(number.one, numero.uno)
537 self.assertEqual(number.two, numero.dos)
538 self.assertEqual(number.three, numero.tres)
539
540 def test_wrong_enum_in_call(self):
541 class Monochrome(Enum):
542 black = 0
543 white = 1
544 class Gender(Enum):
545 male = 0
546 female = 1
547 self.assertRaises(ValueError, Monochrome, Gender.male)
548
549 def test_wrong_enum_in_mixed_call(self):
550 class Monochrome(IntEnum):
551 black = 0
552 white = 1
553 class Gender(Enum):
554 male = 0
555 female = 1
556 self.assertRaises(ValueError, Monochrome, Gender.male)
557
558 def test_mixed_enum_in_call_1(self):
559 class Monochrome(IntEnum):
560 black = 0
561 white = 1
562 class Gender(IntEnum):
563 male = 0
564 female = 1
565 self.assertIs(Monochrome(Gender.female), Monochrome.white)
566
567 def test_mixed_enum_in_call_2(self):
568 class Monochrome(Enum):
569 black = 0
570 white = 1
571 class Gender(IntEnum):
572 male = 0
573 female = 1
574 self.assertIs(Monochrome(Gender.male), Monochrome.black)
575
576 def test_flufl_enum(self):
577 class Fluflnum(Enum):
578 def __int__(self):
579 return int(self.value)
580 class MailManOptions(Fluflnum):
581 option1 = 1
582 option2 = 2
583 option3 = 3
584 self.assertEqual(int(MailManOptions.option1), 1)
585
586 def test_no_such_enum_member(self):
587 class Color(Enum):
588 red = 1
589 green = 2
590 blue = 3
591 with self.assertRaises(ValueError):
592 Color(4)
593 with self.assertRaises(KeyError):
594 Color['chartreuse']
595
596 def test_new_repr(self):
597 class Color(Enum):
598 red = 1
599 green = 2
600 blue = 3
601 def __repr__(self):
602 return "don't you just love shades of %s?" % self.name
603 self.assertEqual(
604 repr(Color.blue),
605 "don't you just love shades of blue?",
606 )
607
608 def test_inherited_repr(self):
609 class MyEnum(Enum):
610 def __repr__(self):
611 return "My name is %s." % self.name
612 class MyIntEnum(int, MyEnum):
613 this = 1
614 that = 2
615 theother = 3
616 self.assertEqual(repr(MyIntEnum.that), "My name is that.")
617
618 def test_multiple_mixin_mro(self):
619 class auto_enum(type(Enum)):
620 def __new__(metacls, cls, bases, classdict):
621 temp = type(classdict)()
622 names = set(classdict._member_names)
623 i = 0
624 for k in classdict._member_names:
625 v = classdict[k]
626 if v is Ellipsis:
627 v = i
628 else:
629 i = v
630 i += 1
631 temp[k] = v
632 for k, v in classdict.items():
633 if k not in names:
634 temp[k] = v
635 return super(auto_enum, metacls).__new__(
636 metacls, cls, bases, temp)
637
638 class AutoNumberedEnum(Enum, metaclass=auto_enum):
639 pass
640
641 class AutoIntEnum(IntEnum, metaclass=auto_enum):
642 pass
643
644 class TestAutoNumber(AutoNumberedEnum):
645 a = ...
646 b = 3
647 c = ...
648
649 class TestAutoInt(AutoIntEnum):
650 a = ...
651 b = 3
652 c = ...
653
654 def test_subclasses_with_getnewargs(self):
655 class NamedInt(int):
656 def __new__(cls, *args):
657 _args = args
658 name, *args = args
659 if len(args) == 0:
660 raise TypeError("name and value must be specified")
661 self = int.__new__(cls, *args)
662 self._intname = name
663 self._args = _args
664 return self
665 def __getnewargs__(self):
666 return self._args
667 @property
668 def __name__(self):
669 return self._intname
670 def __repr__(self):
671 # repr() is updated to include the name and type info
672 return "{}({!r}, {})".format(type(self).__name__,
673 self.__name__,
674 int.__repr__(self))
675 def __str__(self):
676 # str() is unchanged, even if it relies on the repr() fallback
677 base = int
678 base_str = base.__str__
679 if base_str.__objclass__ is object:
680 return base.__repr__(self)
681 return base_str(self)
682 # for simplicity, we only define one operator that
683 # propagates expressions
684 def __add__(self, other):
685 temp = int(self) + int( other)
686 if isinstance(self, NamedInt) and isinstance(other, NamedInt):
687 return NamedInt(
688 '({0} + {1})'.format(self.__name__, other.__name__),
689 temp )
690 else:
691 return temp
692
693 class NEI(NamedInt, Enum):
694 x = ('the-x', 1)
695 y = ('the-y', 2)
696
Ethan Furman2aa27322013-07-19 19:35:56 -0700697
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700698 self.assertIs(NEI.__new__, Enum.__new__)
699 self.assertEqual(repr(NEI.x + NEI.y), "NamedInt('(the-x + the-y)', 3)")
700 globals()['NamedInt'] = NamedInt
701 globals()['NEI'] = NEI
702 NI5 = NamedInt('test', 5)
703 self.assertEqual(NI5, 5)
704 self.assertEqual(loads(dumps(NI5)), 5)
705 self.assertEqual(NEI.y.value, 2)
706 self.assertIs(loads(dumps(NEI.y)), NEI.y)
707
708 def test_subclasses_without_getnewargs(self):
709 class NamedInt(int):
710 def __new__(cls, *args):
711 _args = args
712 name, *args = args
713 if len(args) == 0:
714 raise TypeError("name and value must be specified")
715 self = int.__new__(cls, *args)
716 self._intname = name
717 self._args = _args
718 return self
719 @property
720 def __name__(self):
721 return self._intname
722 def __repr__(self):
723 # repr() is updated to include the name and type info
724 return "{}({!r}, {})".format(type(self).__name__,
725 self.__name__,
726 int.__repr__(self))
727 def __str__(self):
728 # str() is unchanged, even if it relies on the repr() fallback
729 base = int
730 base_str = base.__str__
731 if base_str.__objclass__ is object:
732 return base.__repr__(self)
733 return base_str(self)
734 # for simplicity, we only define one operator that
735 # propagates expressions
736 def __add__(self, other):
737 temp = int(self) + int( other)
738 if isinstance(self, NamedInt) and isinstance(other, NamedInt):
739 return NamedInt(
740 '({0} + {1})'.format(self.__name__, other.__name__),
741 temp )
742 else:
743 return temp
744
745 class NEI(NamedInt, Enum):
746 x = ('the-x', 1)
747 y = ('the-y', 2)
748
749 self.assertIs(NEI.__new__, Enum.__new__)
750 self.assertEqual(repr(NEI.x + NEI.y), "NamedInt('(the-x + the-y)', 3)")
751 globals()['NamedInt'] = NamedInt
752 globals()['NEI'] = NEI
753 NI5 = NamedInt('test', 5)
754 self.assertEqual(NI5, 5)
755 self.assertEqual(NEI.y.value, 2)
756 with self.assertRaises(TypeError):
757 dumps(NEI.x)
758 with self.assertRaises(PicklingError):
759 dumps(NEI)
760
761 def test_tuple_subclass(self):
762 class SomeTuple(tuple, Enum):
763 first = (1, 'for the money')
764 second = (2, 'for the show')
765 third = (3, 'for the music')
766 self.assertIs(type(SomeTuple.first), SomeTuple)
767 self.assertIsInstance(SomeTuple.second, tuple)
768 self.assertEqual(SomeTuple.third, (3, 'for the music'))
769 globals()['SomeTuple'] = SomeTuple
770 self.assertIs(loads(dumps(SomeTuple.first)), SomeTuple.first)
771
772 def test_duplicate_values_give_unique_enum_items(self):
773 class AutoNumber(Enum):
774 first = ()
775 second = ()
776 third = ()
777 def __new__(cls):
778 value = len(cls.__members__) + 1
779 obj = object.__new__(cls)
Ethan Furman520ad572013-07-19 19:47:21 -0700780 obj._value_ = value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700781 return obj
782 def __int__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700783 return int(self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700784 self.assertEqual(
785 list(AutoNumber),
786 [AutoNumber.first, AutoNumber.second, AutoNumber.third],
787 )
788 self.assertEqual(int(AutoNumber.second), 2)
Ethan Furman2aa27322013-07-19 19:35:56 -0700789 self.assertEqual(AutoNumber.third.value, 3)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700790 self.assertIs(AutoNumber(1), AutoNumber.first)
791
792 def test_inherited_new_from_enhanced_enum(self):
793 class AutoNumber(Enum):
794 def __new__(cls):
795 value = len(cls.__members__) + 1
796 obj = object.__new__(cls)
Ethan Furman520ad572013-07-19 19:47:21 -0700797 obj._value_ = value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700798 return obj
799 def __int__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700800 return int(self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700801 class Color(AutoNumber):
802 red = ()
803 green = ()
804 blue = ()
805 self.assertEqual(list(Color), [Color.red, Color.green, Color.blue])
806 self.assertEqual(list(map(int, Color)), [1, 2, 3])
807
808 def test_inherited_new_from_mixed_enum(self):
809 class AutoNumber(IntEnum):
810 def __new__(cls):
811 value = len(cls.__members__) + 1
812 obj = int.__new__(cls, value)
Ethan Furman520ad572013-07-19 19:47:21 -0700813 obj._value_ = value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700814 return obj
815 class Color(AutoNumber):
816 red = ()
817 green = ()
818 blue = ()
819 self.assertEqual(list(Color), [Color.red, Color.green, Color.blue])
820 self.assertEqual(list(map(int, Color)), [1, 2, 3])
821
822 def test_ordered_mixin(self):
823 class OrderedEnum(Enum):
824 def __ge__(self, other):
825 if self.__class__ is other.__class__:
Ethan Furman520ad572013-07-19 19:47:21 -0700826 return self._value_ >= other._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700827 return NotImplemented
828 def __gt__(self, other):
829 if self.__class__ is other.__class__:
Ethan Furman520ad572013-07-19 19:47:21 -0700830 return self._value_ > other._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700831 return NotImplemented
832 def __le__(self, other):
833 if self.__class__ is other.__class__:
Ethan Furman520ad572013-07-19 19:47:21 -0700834 return self._value_ <= other._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700835 return NotImplemented
836 def __lt__(self, other):
837 if self.__class__ is other.__class__:
Ethan Furman520ad572013-07-19 19:47:21 -0700838 return self._value_ < other._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700839 return NotImplemented
840 class Grade(OrderedEnum):
841 A = 5
842 B = 4
843 C = 3
844 D = 2
845 F = 1
846 self.assertGreater(Grade.A, Grade.B)
847 self.assertLessEqual(Grade.F, Grade.C)
848 self.assertLess(Grade.D, Grade.A)
849 self.assertGreaterEqual(Grade.B, Grade.B)
Ethan Furman520ad572013-07-19 19:47:21 -0700850
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700851 def test_extending2(self):
852 class Shade(Enum):
853 def shade(self):
854 print(self.name)
855 class Color(Shade):
856 red = 1
857 green = 2
858 blue = 3
859 with self.assertRaises(TypeError):
860 class MoreColor(Color):
861 cyan = 4
862 magenta = 5
863 yellow = 6
864
865 def test_extending3(self):
866 class Shade(Enum):
867 def shade(self):
868 return self.name
869 class Color(Shade):
870 def hex(self):
871 return '%s hexlified!' % self.value
872 class MoreColor(Color):
873 cyan = 4
874 magenta = 5
875 yellow = 6
876 self.assertEqual(MoreColor.magenta.hex(), '5 hexlified!')
877
878
879 def test_no_duplicates(self):
880 class UniqueEnum(Enum):
881 def __init__(self, *args):
882 cls = self.__class__
883 if any(self.value == e.value for e in cls):
884 a = self.name
885 e = cls(self.value).name
886 raise ValueError(
887 "aliases not allowed in UniqueEnum: %r --> %r"
888 % (a, e)
889 )
890 class Color(UniqueEnum):
891 red = 1
892 green = 2
893 blue = 3
894 with self.assertRaises(ValueError):
895 class Color(UniqueEnum):
896 red = 1
897 green = 2
898 blue = 3
899 grene = 2
900
901 def test_init(self):
902 class Planet(Enum):
903 MERCURY = (3.303e+23, 2.4397e6)
904 VENUS = (4.869e+24, 6.0518e6)
905 EARTH = (5.976e+24, 6.37814e6)
906 MARS = (6.421e+23, 3.3972e6)
907 JUPITER = (1.9e+27, 7.1492e7)
908 SATURN = (5.688e+26, 6.0268e7)
909 URANUS = (8.686e+25, 2.5559e7)
910 NEPTUNE = (1.024e+26, 2.4746e7)
911 def __init__(self, mass, radius):
912 self.mass = mass # in kilograms
913 self.radius = radius # in meters
914 @property
915 def surface_gravity(self):
916 # universal gravitational constant (m3 kg-1 s-2)
917 G = 6.67300E-11
918 return G * self.mass / (self.radius * self.radius)
919 self.assertEqual(round(Planet.EARTH.surface_gravity, 2), 9.80)
920 self.assertEqual(Planet.EARTH.value, (5.976e+24, 6.37814e6))
921
Ethan Furman2aa27322013-07-19 19:35:56 -0700922 def test_nonhash_value(self):
923 class AutoNumberInAList(Enum):
924 def __new__(cls):
925 value = [len(cls.__members__) + 1]
926 obj = object.__new__(cls)
Ethan Furman520ad572013-07-19 19:47:21 -0700927 obj._value_ = value
Ethan Furman2aa27322013-07-19 19:35:56 -0700928 return obj
929 class ColorInAList(AutoNumberInAList):
930 red = ()
931 green = ()
932 blue = ()
933 self.assertEqual(list(ColorInAList), [ColorInAList.red, ColorInAList.green, ColorInAList.blue])
934 self.assertEqual(ColorInAList.red.value, [1])
935 self.assertEqual(ColorInAList([1]), ColorInAList.red)
936
937
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700938
Ethan Furmanf24bb352013-07-18 17:05:39 -0700939class TestUnique(unittest.TestCase):
940
941 def test_unique_clean(self):
942 @unique
943 class Clean(Enum):
944 one = 1
945 two = 'dos'
946 tres = 4.0
947 @unique
948 class Cleaner(IntEnum):
949 single = 1
950 double = 2
951 triple = 3
952
953 def test_unique_dirty(self):
954 with self.assertRaisesRegex(ValueError, 'tres.*one'):
955 @unique
956 class Dirty(Enum):
957 one = 1
958 two = 'dos'
959 tres = 1
960 with self.assertRaisesRegex(
961 ValueError,
962 'double.*single.*turkey.*triple',
963 ):
964 @unique
965 class Dirtier(IntEnum):
966 single = 1
967 double = 1
968 triple = 3
969 turkey = 3
970
971
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700972if __name__ == '__main__':
973 unittest.main()