blob: d59c5e3ea3cc89871e7b78c6c504a8beccfc016d [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
Ethan Furmanec15a822013-08-31 19:17:41 -070070 class Konstants(float, Enum):
71 E = 2.7182818
72 PI = 3.1415926
73 TAU = 2 * PI
74 self.Konstants = Konstants
75
76 class Grades(IntEnum):
77 A = 5
78 B = 4
79 C = 3
80 D = 2
81 F = 0
82 self.Grades = Grades
83
84 class Directional(str, Enum):
85 EAST = 'east'
86 WEST = 'west'
87 NORTH = 'north'
88 SOUTH = 'south'
89 self.Directional = Directional
90
91 from datetime import date
92 class Holiday(date, Enum):
93 NEW_YEAR = 2013, 1, 1
94 IDES_OF_MARCH = 2013, 3, 15
95 self.Holiday = Holiday
96
Ethan Furman388a3922013-08-12 06:51:41 -070097 def test_dir_on_class(self):
98 Season = self.Season
99 self.assertEqual(
100 set(dir(Season)),
Ethan Furmanc850f342013-09-15 16:59:35 -0700101 set(['__class__', '__doc__', '__members__', '__module__',
Ethan Furman388a3922013-08-12 06:51:41 -0700102 'SPRING', 'SUMMER', 'AUTUMN', 'WINTER']),
103 )
104
105 def test_dir_on_item(self):
106 Season = self.Season
107 self.assertEqual(
108 set(dir(Season.WINTER)),
Ethan Furmanc850f342013-09-15 16:59:35 -0700109 set(['__class__', '__doc__', '__module__', 'name', 'value']),
Ethan Furman388a3922013-08-12 06:51:41 -0700110 )
111
Ethan Furmanc850f342013-09-15 16:59:35 -0700112 def test_dir_with_added_behavior(self):
113 class Test(Enum):
114 this = 'that'
115 these = 'those'
116 def wowser(self):
117 return ("Wowser! I'm %s!" % self.name)
118 self.assertEqual(
119 set(dir(Test)),
120 set(['__class__', '__doc__', '__members__', '__module__', 'this', 'these']),
121 )
122 self.assertEqual(
123 set(dir(Test.this)),
124 set(['__class__', '__doc__', '__module__', 'name', 'value', 'wowser']),
125 )
126
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700127 def test_enum_in_enum_out(self):
128 Season = self.Season
129 self.assertIs(Season(Season.WINTER), Season.WINTER)
130
131 def test_enum_value(self):
132 Season = self.Season
133 self.assertEqual(Season.SPRING.value, 1)
134
135 def test_intenum_value(self):
136 self.assertEqual(IntStooges.CURLY.value, 2)
137
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700138 def test_enum(self):
139 Season = self.Season
140 lst = list(Season)
141 self.assertEqual(len(lst), len(Season))
142 self.assertEqual(len(Season), 4, Season)
143 self.assertEqual(
144 [Season.SPRING, Season.SUMMER, Season.AUTUMN, Season.WINTER], lst)
145
146 for i, season in enumerate('SPRING SUMMER AUTUMN WINTER'.split(), 1):
147 e = Season(i)
148 self.assertEqual(e, getattr(Season, season))
149 self.assertEqual(e.value, i)
150 self.assertNotEqual(e, i)
151 self.assertEqual(e.name, season)
152 self.assertIn(e, Season)
153 self.assertIs(type(e), Season)
154 self.assertIsInstance(e, Season)
155 self.assertEqual(str(e), 'Season.' + season)
156 self.assertEqual(
157 repr(e),
158 '<Season.{0}: {1}>'.format(season, i),
159 )
160
161 def test_value_name(self):
162 Season = self.Season
163 self.assertEqual(Season.SPRING.name, 'SPRING')
164 self.assertEqual(Season.SPRING.value, 1)
165 with self.assertRaises(AttributeError):
166 Season.SPRING.name = 'invierno'
167 with self.assertRaises(AttributeError):
168 Season.SPRING.value = 2
169
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700170 def test_changing_member(self):
171 Season = self.Season
172 with self.assertRaises(AttributeError):
173 Season.WINTER = 'really cold'
174
Ethan Furman64a99722013-09-22 16:18:19 -0700175 def test_attribute_deletion(self):
176 class Season(Enum):
177 SPRING = 1
178 SUMMER = 2
179 AUTUMN = 3
180 WINTER = 4
181
182 def spam(cls):
183 pass
184
185 self.assertTrue(hasattr(Season, 'spam'))
186 del Season.spam
187 self.assertFalse(hasattr(Season, 'spam'))
188
189 with self.assertRaises(AttributeError):
190 del Season.SPRING
191 with self.assertRaises(AttributeError):
192 del Season.DRY
193 with self.assertRaises(AttributeError):
194 del Season.SPRING.name
195
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700196 def test_invalid_names(self):
197 with self.assertRaises(ValueError):
198 class Wrong(Enum):
199 mro = 9
200 with self.assertRaises(ValueError):
201 class Wrong(Enum):
202 _create_= 11
203 with self.assertRaises(ValueError):
204 class Wrong(Enum):
205 _get_mixins_ = 9
206 with self.assertRaises(ValueError):
207 class Wrong(Enum):
208 _find_new_ = 1
209 with self.assertRaises(ValueError):
210 class Wrong(Enum):
211 _any_name_ = 9
212
213 def test_contains(self):
214 Season = self.Season
215 self.assertIn(Season.AUTUMN, Season)
216 self.assertNotIn(3, Season)
217
218 val = Season(3)
219 self.assertIn(val, Season)
220
221 class OtherEnum(Enum):
222 one = 1; two = 2
223 self.assertNotIn(OtherEnum.two, Season)
224
225 def test_comparisons(self):
226 Season = self.Season
227 with self.assertRaises(TypeError):
228 Season.SPRING < Season.WINTER
229 with self.assertRaises(TypeError):
230 Season.SPRING > 4
231
232 self.assertNotEqual(Season.SPRING, 1)
233
234 class Part(Enum):
235 SPRING = 1
236 CLIP = 2
237 BARREL = 3
238
239 self.assertNotEqual(Season.SPRING, Part.SPRING)
240 with self.assertRaises(TypeError):
241 Season.SPRING < Part.CLIP
242
243 def test_enum_duplicates(self):
244 class Season(Enum):
245 SPRING = 1
246 SUMMER = 2
247 AUTUMN = FALL = 3
248 WINTER = 4
249 ANOTHER_SPRING = 1
250 lst = list(Season)
251 self.assertEqual(
252 lst,
253 [Season.SPRING, Season.SUMMER,
254 Season.AUTUMN, Season.WINTER,
255 ])
256 self.assertIs(Season.FALL, Season.AUTUMN)
257 self.assertEqual(Season.FALL.value, 3)
258 self.assertEqual(Season.AUTUMN.value, 3)
259 self.assertIs(Season(3), Season.AUTUMN)
260 self.assertIs(Season(1), Season.SPRING)
261 self.assertEqual(Season.FALL.name, 'AUTUMN')
262 self.assertEqual(
263 [k for k,v in Season.__members__.items() if v.name != k],
264 ['FALL', 'ANOTHER_SPRING'],
265 )
266
Ethan Furman101e0742013-09-15 12:34:36 -0700267 def test_duplicate_name(self):
268 with self.assertRaises(TypeError):
269 class Color(Enum):
270 red = 1
271 green = 2
272 blue = 3
273 red = 4
274
275 with self.assertRaises(TypeError):
276 class Color(Enum):
277 red = 1
278 green = 2
279 blue = 3
280 def red(self):
281 return 'red'
282
283 with self.assertRaises(TypeError):
284 class Color(Enum):
285 @property
286 def red(self):
287 return 'redder'
288 red = 1
289 green = 2
290 blue = 3
291
292
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700293 def test_enum_with_value_name(self):
294 class Huh(Enum):
295 name = 1
296 value = 2
297 self.assertEqual(
298 list(Huh),
299 [Huh.name, Huh.value],
300 )
301 self.assertIs(type(Huh.name), Huh)
302 self.assertEqual(Huh.name.name, 'name')
303 self.assertEqual(Huh.name.value, 1)
Ethan Furmanec15a822013-08-31 19:17:41 -0700304
305 def test_format_enum(self):
306 Season = self.Season
307 self.assertEqual('{}'.format(Season.SPRING),
308 '{}'.format(str(Season.SPRING)))
309 self.assertEqual( '{:}'.format(Season.SPRING),
310 '{:}'.format(str(Season.SPRING)))
311 self.assertEqual('{:20}'.format(Season.SPRING),
312 '{:20}'.format(str(Season.SPRING)))
313 self.assertEqual('{:^20}'.format(Season.SPRING),
314 '{:^20}'.format(str(Season.SPRING)))
315 self.assertEqual('{:>20}'.format(Season.SPRING),
316 '{:>20}'.format(str(Season.SPRING)))
317 self.assertEqual('{:<20}'.format(Season.SPRING),
318 '{:<20}'.format(str(Season.SPRING)))
319
320 def test_format_enum_custom(self):
321 class TestFloat(float, Enum):
322 one = 1.0
323 two = 2.0
324 def __format__(self, spec):
325 return 'TestFloat success!'
326 self.assertEqual('{}'.format(TestFloat.one), 'TestFloat success!')
327
328 def assertFormatIsValue(self, spec, member):
329 self.assertEqual(spec.format(member), spec.format(member.value))
330
331 def test_format_enum_date(self):
332 Holiday = self.Holiday
333 self.assertFormatIsValue('{}', Holiday.IDES_OF_MARCH)
334 self.assertFormatIsValue('{:}', Holiday.IDES_OF_MARCH)
335 self.assertFormatIsValue('{:20}', Holiday.IDES_OF_MARCH)
336 self.assertFormatIsValue('{:^20}', Holiday.IDES_OF_MARCH)
337 self.assertFormatIsValue('{:>20}', Holiday.IDES_OF_MARCH)
338 self.assertFormatIsValue('{:<20}', Holiday.IDES_OF_MARCH)
339 self.assertFormatIsValue('{:%Y %m}', Holiday.IDES_OF_MARCH)
340 self.assertFormatIsValue('{:%Y %m %M:00}', Holiday.IDES_OF_MARCH)
341
342 def test_format_enum_float(self):
343 Konstants = self.Konstants
344 self.assertFormatIsValue('{}', Konstants.TAU)
345 self.assertFormatIsValue('{:}', Konstants.TAU)
346 self.assertFormatIsValue('{:20}', Konstants.TAU)
347 self.assertFormatIsValue('{:^20}', Konstants.TAU)
348 self.assertFormatIsValue('{:>20}', Konstants.TAU)
349 self.assertFormatIsValue('{:<20}', Konstants.TAU)
350 self.assertFormatIsValue('{:n}', Konstants.TAU)
351 self.assertFormatIsValue('{:5.2}', Konstants.TAU)
352 self.assertFormatIsValue('{:f}', Konstants.TAU)
353
354 def test_format_enum_int(self):
355 Grades = self.Grades
356 self.assertFormatIsValue('{}', Grades.C)
357 self.assertFormatIsValue('{:}', Grades.C)
358 self.assertFormatIsValue('{:20}', Grades.C)
359 self.assertFormatIsValue('{:^20}', Grades.C)
360 self.assertFormatIsValue('{:>20}', Grades.C)
361 self.assertFormatIsValue('{:<20}', Grades.C)
362 self.assertFormatIsValue('{:+}', Grades.C)
363 self.assertFormatIsValue('{:08X}', Grades.C)
364 self.assertFormatIsValue('{:b}', Grades.C)
365
366 def test_format_enum_str(self):
367 Directional = self.Directional
368 self.assertFormatIsValue('{}', Directional.WEST)
369 self.assertFormatIsValue('{:}', Directional.WEST)
370 self.assertFormatIsValue('{:20}', Directional.WEST)
371 self.assertFormatIsValue('{:^20}', Directional.WEST)
372 self.assertFormatIsValue('{:>20}', Directional.WEST)
373 self.assertFormatIsValue('{:<20}', Directional.WEST)
374
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700375 def test_hash(self):
376 Season = self.Season
377 dates = {}
378 dates[Season.WINTER] = '1225'
379 dates[Season.SPRING] = '0315'
380 dates[Season.SUMMER] = '0704'
381 dates[Season.AUTUMN] = '1031'
382 self.assertEqual(dates[Season.AUTUMN], '1031')
383
384 def test_intenum_from_scratch(self):
385 class phy(int, Enum):
386 pi = 3
387 tau = 2 * pi
388 self.assertTrue(phy.pi < phy.tau)
389
390 def test_intenum_inherited(self):
391 class IntEnum(int, Enum):
392 pass
393 class phy(IntEnum):
394 pi = 3
395 tau = 2 * pi
396 self.assertTrue(phy.pi < phy.tau)
397
398 def test_floatenum_from_scratch(self):
399 class phy(float, Enum):
Ethan Furmanec15a822013-08-31 19:17:41 -0700400 pi = 3.1415926
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700401 tau = 2 * pi
402 self.assertTrue(phy.pi < phy.tau)
403
404 def test_floatenum_inherited(self):
405 class FloatEnum(float, Enum):
406 pass
407 class phy(FloatEnum):
Ethan Furmanec15a822013-08-31 19:17:41 -0700408 pi = 3.1415926
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700409 tau = 2 * pi
410 self.assertTrue(phy.pi < phy.tau)
411
412 def test_strenum_from_scratch(self):
413 class phy(str, Enum):
414 pi = 'Pi'
415 tau = 'Tau'
416 self.assertTrue(phy.pi < phy.tau)
417
418 def test_strenum_inherited(self):
419 class StrEnum(str, Enum):
420 pass
421 class phy(StrEnum):
422 pi = 'Pi'
423 tau = 'Tau'
424 self.assertTrue(phy.pi < phy.tau)
425
426
427 def test_intenum(self):
428 class WeekDay(IntEnum):
429 SUNDAY = 1
430 MONDAY = 2
431 TUESDAY = 3
432 WEDNESDAY = 4
433 THURSDAY = 5
434 FRIDAY = 6
435 SATURDAY = 7
436
437 self.assertEqual(['a', 'b', 'c'][WeekDay.MONDAY], 'c')
438 self.assertEqual([i for i in range(WeekDay.TUESDAY)], [0, 1, 2])
439
440 lst = list(WeekDay)
441 self.assertEqual(len(lst), len(WeekDay))
442 self.assertEqual(len(WeekDay), 7)
443 target = 'SUNDAY MONDAY TUESDAY WEDNESDAY THURSDAY FRIDAY SATURDAY'
444 target = target.split()
445 for i, weekday in enumerate(target, 1):
446 e = WeekDay(i)
447 self.assertEqual(e, i)
448 self.assertEqual(int(e), i)
449 self.assertEqual(e.name, weekday)
450 self.assertIn(e, WeekDay)
451 self.assertEqual(lst.index(e)+1, i)
452 self.assertTrue(0 < e < 8)
453 self.assertIs(type(e), WeekDay)
454 self.assertIsInstance(e, int)
455 self.assertIsInstance(e, Enum)
456
457 def test_intenum_duplicates(self):
458 class WeekDay(IntEnum):
459 SUNDAY = 1
460 MONDAY = 2
461 TUESDAY = TEUSDAY = 3
462 WEDNESDAY = 4
463 THURSDAY = 5
464 FRIDAY = 6
465 SATURDAY = 7
466 self.assertIs(WeekDay.TEUSDAY, WeekDay.TUESDAY)
467 self.assertEqual(WeekDay(3).name, 'TUESDAY')
468 self.assertEqual([k for k,v in WeekDay.__members__.items()
469 if v.name != k], ['TEUSDAY', ])
470
471 def test_pickle_enum(self):
472 if isinstance(Stooges, Exception):
473 raise Stooges
474 self.assertIs(Stooges.CURLY, loads(dumps(Stooges.CURLY)))
475 self.assertIs(Stooges, loads(dumps(Stooges)))
476
477 def test_pickle_int(self):
478 if isinstance(IntStooges, Exception):
479 raise IntStooges
480 self.assertIs(IntStooges.CURLY, loads(dumps(IntStooges.CURLY)))
481 self.assertIs(IntStooges, loads(dumps(IntStooges)))
482
483 def test_pickle_float(self):
484 if isinstance(FloatStooges, Exception):
485 raise FloatStooges
486 self.assertIs(FloatStooges.CURLY, loads(dumps(FloatStooges.CURLY)))
487 self.assertIs(FloatStooges, loads(dumps(FloatStooges)))
488
489 def test_pickle_enum_function(self):
490 if isinstance(Answer, Exception):
491 raise Answer
492 self.assertIs(Answer.him, loads(dumps(Answer.him)))
493 self.assertIs(Answer, loads(dumps(Answer)))
494
495 def test_pickle_enum_function_with_module(self):
496 if isinstance(Question, Exception):
497 raise Question
498 self.assertIs(Question.who, loads(dumps(Question.who)))
499 self.assertIs(Question, loads(dumps(Question)))
500
501 def test_exploding_pickle(self):
502 BadPickle = Enum('BadPickle', 'dill sweet bread-n-butter')
503 enum._make_class_unpicklable(BadPickle)
504 globals()['BadPickle'] = BadPickle
505 with self.assertRaises(TypeError):
506 dumps(BadPickle.dill)
507 with self.assertRaises(PicklingError):
508 dumps(BadPickle)
509
510 def test_string_enum(self):
511 class SkillLevel(str, Enum):
512 master = 'what is the sound of one hand clapping?'
513 journeyman = 'why did the chicken cross the road?'
514 apprentice = 'knock, knock!'
515 self.assertEqual(SkillLevel.apprentice, 'knock, knock!')
516
517 def test_getattr_getitem(self):
518 class Period(Enum):
519 morning = 1
520 noon = 2
521 evening = 3
522 night = 4
523 self.assertIs(Period(2), Period.noon)
524 self.assertIs(getattr(Period, 'night'), Period.night)
525 self.assertIs(Period['morning'], Period.morning)
526
527 def test_getattr_dunder(self):
528 Season = self.Season
529 self.assertTrue(getattr(Season, '__eq__'))
530
531 def test_iteration_order(self):
532 class Season(Enum):
533 SUMMER = 2
534 WINTER = 4
535 AUTUMN = 3
536 SPRING = 1
537 self.assertEqual(
538 list(Season),
539 [Season.SUMMER, Season.WINTER, Season.AUTUMN, Season.SPRING],
540 )
541
Ethan Furman2131a4a2013-09-14 18:11:24 -0700542 def test_reversed_iteration_order(self):
543 self.assertEqual(
544 list(reversed(self.Season)),
545 [self.Season.WINTER, self.Season.AUTUMN, self.Season.SUMMER,
546 self.Season.SPRING]
547 )
548
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700549 def test_programatic_function_string(self):
550 SummerMonth = Enum('SummerMonth', 'june july august')
551 lst = list(SummerMonth)
552 self.assertEqual(len(lst), len(SummerMonth))
553 self.assertEqual(len(SummerMonth), 3, SummerMonth)
554 self.assertEqual(
555 [SummerMonth.june, SummerMonth.july, SummerMonth.august],
556 lst,
557 )
558 for i, month in enumerate('june july august'.split(), 1):
559 e = SummerMonth(i)
560 self.assertEqual(int(e.value), i)
561 self.assertNotEqual(e, i)
562 self.assertEqual(e.name, month)
563 self.assertIn(e, SummerMonth)
564 self.assertIs(type(e), SummerMonth)
565
566 def test_programatic_function_string_list(self):
567 SummerMonth = Enum('SummerMonth', ['june', 'july', 'august'])
568 lst = list(SummerMonth)
569 self.assertEqual(len(lst), len(SummerMonth))
570 self.assertEqual(len(SummerMonth), 3, SummerMonth)
571 self.assertEqual(
572 [SummerMonth.june, SummerMonth.july, SummerMonth.august],
573 lst,
574 )
575 for i, month in enumerate('june july august'.split(), 1):
576 e = SummerMonth(i)
577 self.assertEqual(int(e.value), i)
578 self.assertNotEqual(e, i)
579 self.assertEqual(e.name, month)
580 self.assertIn(e, SummerMonth)
581 self.assertIs(type(e), SummerMonth)
582
583 def test_programatic_function_iterable(self):
584 SummerMonth = Enum(
585 'SummerMonth',
586 (('june', 1), ('july', 2), ('august', 3))
587 )
588 lst = list(SummerMonth)
589 self.assertEqual(len(lst), len(SummerMonth))
590 self.assertEqual(len(SummerMonth), 3, SummerMonth)
591 self.assertEqual(
592 [SummerMonth.june, SummerMonth.july, SummerMonth.august],
593 lst,
594 )
595 for i, month in enumerate('june july august'.split(), 1):
596 e = SummerMonth(i)
597 self.assertEqual(int(e.value), i)
598 self.assertNotEqual(e, i)
599 self.assertEqual(e.name, month)
600 self.assertIn(e, SummerMonth)
601 self.assertIs(type(e), SummerMonth)
602
603 def test_programatic_function_from_dict(self):
604 SummerMonth = Enum(
605 'SummerMonth',
606 OrderedDict((('june', 1), ('july', 2), ('august', 3)))
607 )
608 lst = list(SummerMonth)
609 self.assertEqual(len(lst), len(SummerMonth))
610 self.assertEqual(len(SummerMonth), 3, SummerMonth)
611 self.assertEqual(
612 [SummerMonth.june, SummerMonth.july, SummerMonth.august],
613 lst,
614 )
615 for i, month in enumerate('june july august'.split(), 1):
616 e = SummerMonth(i)
617 self.assertEqual(int(e.value), i)
618 self.assertNotEqual(e, i)
619 self.assertEqual(e.name, month)
620 self.assertIn(e, SummerMonth)
621 self.assertIs(type(e), SummerMonth)
622
623 def test_programatic_function_type(self):
624 SummerMonth = Enum('SummerMonth', 'june july august', type=int)
625 lst = list(SummerMonth)
626 self.assertEqual(len(lst), len(SummerMonth))
627 self.assertEqual(len(SummerMonth), 3, SummerMonth)
628 self.assertEqual(
629 [SummerMonth.june, SummerMonth.july, SummerMonth.august],
630 lst,
631 )
632 for i, month in enumerate('june july august'.split(), 1):
633 e = SummerMonth(i)
634 self.assertEqual(e, i)
635 self.assertEqual(e.name, month)
636 self.assertIn(e, SummerMonth)
637 self.assertIs(type(e), SummerMonth)
638
639 def test_programatic_function_type_from_subclass(self):
640 SummerMonth = IntEnum('SummerMonth', 'june july august')
641 lst = list(SummerMonth)
642 self.assertEqual(len(lst), len(SummerMonth))
643 self.assertEqual(len(SummerMonth), 3, SummerMonth)
644 self.assertEqual(
645 [SummerMonth.june, SummerMonth.july, SummerMonth.august],
646 lst,
647 )
648 for i, month in enumerate('june july august'.split(), 1):
649 e = SummerMonth(i)
650 self.assertEqual(e, i)
651 self.assertEqual(e.name, month)
652 self.assertIn(e, SummerMonth)
653 self.assertIs(type(e), SummerMonth)
654
655 def test_subclassing(self):
656 if isinstance(Name, Exception):
657 raise Name
658 self.assertEqual(Name.BDFL, 'Guido van Rossum')
659 self.assertTrue(Name.BDFL, Name('Guido van Rossum'))
660 self.assertIs(Name.BDFL, getattr(Name, 'BDFL'))
661 self.assertIs(Name.BDFL, loads(dumps(Name.BDFL)))
662
663 def test_extending(self):
664 class Color(Enum):
665 red = 1
666 green = 2
667 blue = 3
668 with self.assertRaises(TypeError):
669 class MoreColor(Color):
670 cyan = 4
671 magenta = 5
672 yellow = 6
673
674 def test_exclude_methods(self):
675 class whatever(Enum):
676 this = 'that'
677 these = 'those'
678 def really(self):
679 return 'no, not %s' % self.value
680 self.assertIsNot(type(whatever.really), whatever)
681 self.assertEqual(whatever.this.really(), 'no, not that')
682
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700683 def test_wrong_inheritance_order(self):
684 with self.assertRaises(TypeError):
685 class Wrong(Enum, str):
686 NotHere = 'error before this point'
687
688 def test_intenum_transitivity(self):
689 class number(IntEnum):
690 one = 1
691 two = 2
692 three = 3
693 class numero(IntEnum):
694 uno = 1
695 dos = 2
696 tres = 3
697 self.assertEqual(number.one, numero.uno)
698 self.assertEqual(number.two, numero.dos)
699 self.assertEqual(number.three, numero.tres)
700
701 def test_wrong_enum_in_call(self):
702 class Monochrome(Enum):
703 black = 0
704 white = 1
705 class Gender(Enum):
706 male = 0
707 female = 1
708 self.assertRaises(ValueError, Monochrome, Gender.male)
709
710 def test_wrong_enum_in_mixed_call(self):
711 class Monochrome(IntEnum):
712 black = 0
713 white = 1
714 class Gender(Enum):
715 male = 0
716 female = 1
717 self.assertRaises(ValueError, Monochrome, Gender.male)
718
719 def test_mixed_enum_in_call_1(self):
720 class Monochrome(IntEnum):
721 black = 0
722 white = 1
723 class Gender(IntEnum):
724 male = 0
725 female = 1
726 self.assertIs(Monochrome(Gender.female), Monochrome.white)
727
728 def test_mixed_enum_in_call_2(self):
729 class Monochrome(Enum):
730 black = 0
731 white = 1
732 class Gender(IntEnum):
733 male = 0
734 female = 1
735 self.assertIs(Monochrome(Gender.male), Monochrome.black)
736
737 def test_flufl_enum(self):
738 class Fluflnum(Enum):
739 def __int__(self):
740 return int(self.value)
741 class MailManOptions(Fluflnum):
742 option1 = 1
743 option2 = 2
744 option3 = 3
745 self.assertEqual(int(MailManOptions.option1), 1)
746
Ethan Furman5e5a8232013-08-04 08:42:23 -0700747 def test_introspection(self):
748 class Number(IntEnum):
749 one = 100
750 two = 200
751 self.assertIs(Number.one._member_type_, int)
752 self.assertIs(Number._member_type_, int)
753 class String(str, Enum):
754 yarn = 'soft'
755 rope = 'rough'
756 wire = 'hard'
757 self.assertIs(String.yarn._member_type_, str)
758 self.assertIs(String._member_type_, str)
759 class Plain(Enum):
760 vanilla = 'white'
761 one = 1
762 self.assertIs(Plain.vanilla._member_type_, object)
763 self.assertIs(Plain._member_type_, object)
764
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700765 def test_no_such_enum_member(self):
766 class Color(Enum):
767 red = 1
768 green = 2
769 blue = 3
770 with self.assertRaises(ValueError):
771 Color(4)
772 with self.assertRaises(KeyError):
773 Color['chartreuse']
774
775 def test_new_repr(self):
776 class Color(Enum):
777 red = 1
778 green = 2
779 blue = 3
780 def __repr__(self):
781 return "don't you just love shades of %s?" % self.name
782 self.assertEqual(
783 repr(Color.blue),
784 "don't you just love shades of blue?",
785 )
786
787 def test_inherited_repr(self):
788 class MyEnum(Enum):
789 def __repr__(self):
790 return "My name is %s." % self.name
791 class MyIntEnum(int, MyEnum):
792 this = 1
793 that = 2
794 theother = 3
795 self.assertEqual(repr(MyIntEnum.that), "My name is that.")
796
797 def test_multiple_mixin_mro(self):
798 class auto_enum(type(Enum)):
799 def __new__(metacls, cls, bases, classdict):
800 temp = type(classdict)()
801 names = set(classdict._member_names)
802 i = 0
803 for k in classdict._member_names:
804 v = classdict[k]
805 if v is Ellipsis:
806 v = i
807 else:
808 i = v
809 i += 1
810 temp[k] = v
811 for k, v in classdict.items():
812 if k not in names:
813 temp[k] = v
814 return super(auto_enum, metacls).__new__(
815 metacls, cls, bases, temp)
816
817 class AutoNumberedEnum(Enum, metaclass=auto_enum):
818 pass
819
820 class AutoIntEnum(IntEnum, metaclass=auto_enum):
821 pass
822
823 class TestAutoNumber(AutoNumberedEnum):
824 a = ...
825 b = 3
826 c = ...
827
828 class TestAutoInt(AutoIntEnum):
829 a = ...
830 b = 3
831 c = ...
832
833 def test_subclasses_with_getnewargs(self):
834 class NamedInt(int):
835 def __new__(cls, *args):
836 _args = args
837 name, *args = args
838 if len(args) == 0:
839 raise TypeError("name and value must be specified")
840 self = int.__new__(cls, *args)
841 self._intname = name
842 self._args = _args
843 return self
844 def __getnewargs__(self):
845 return self._args
846 @property
847 def __name__(self):
848 return self._intname
849 def __repr__(self):
850 # repr() is updated to include the name and type info
851 return "{}({!r}, {})".format(type(self).__name__,
852 self.__name__,
853 int.__repr__(self))
854 def __str__(self):
855 # str() is unchanged, even if it relies on the repr() fallback
856 base = int
857 base_str = base.__str__
858 if base_str.__objclass__ is object:
859 return base.__repr__(self)
860 return base_str(self)
861 # for simplicity, we only define one operator that
862 # propagates expressions
863 def __add__(self, other):
864 temp = int(self) + int( other)
865 if isinstance(self, NamedInt) and isinstance(other, NamedInt):
866 return NamedInt(
867 '({0} + {1})'.format(self.__name__, other.__name__),
868 temp )
869 else:
870 return temp
871
872 class NEI(NamedInt, Enum):
873 x = ('the-x', 1)
874 y = ('the-y', 2)
875
Ethan Furman2aa27322013-07-19 19:35:56 -0700876
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700877 self.assertIs(NEI.__new__, Enum.__new__)
878 self.assertEqual(repr(NEI.x + NEI.y), "NamedInt('(the-x + the-y)', 3)")
879 globals()['NamedInt'] = NamedInt
880 globals()['NEI'] = NEI
881 NI5 = NamedInt('test', 5)
882 self.assertEqual(NI5, 5)
883 self.assertEqual(loads(dumps(NI5)), 5)
884 self.assertEqual(NEI.y.value, 2)
885 self.assertIs(loads(dumps(NEI.y)), NEI.y)
886
887 def test_subclasses_without_getnewargs(self):
888 class NamedInt(int):
889 def __new__(cls, *args):
890 _args = args
891 name, *args = args
892 if len(args) == 0:
893 raise TypeError("name and value must be specified")
894 self = int.__new__(cls, *args)
895 self._intname = name
896 self._args = _args
897 return self
898 @property
899 def __name__(self):
900 return self._intname
901 def __repr__(self):
902 # repr() is updated to include the name and type info
903 return "{}({!r}, {})".format(type(self).__name__,
904 self.__name__,
905 int.__repr__(self))
906 def __str__(self):
907 # str() is unchanged, even if it relies on the repr() fallback
908 base = int
909 base_str = base.__str__
910 if base_str.__objclass__ is object:
911 return base.__repr__(self)
912 return base_str(self)
913 # for simplicity, we only define one operator that
914 # propagates expressions
915 def __add__(self, other):
916 temp = int(self) + int( other)
917 if isinstance(self, NamedInt) and isinstance(other, NamedInt):
918 return NamedInt(
919 '({0} + {1})'.format(self.__name__, other.__name__),
920 temp )
921 else:
922 return temp
923
924 class NEI(NamedInt, Enum):
925 x = ('the-x', 1)
926 y = ('the-y', 2)
927
928 self.assertIs(NEI.__new__, Enum.__new__)
929 self.assertEqual(repr(NEI.x + NEI.y), "NamedInt('(the-x + the-y)', 3)")
930 globals()['NamedInt'] = NamedInt
931 globals()['NEI'] = NEI
932 NI5 = NamedInt('test', 5)
933 self.assertEqual(NI5, 5)
934 self.assertEqual(NEI.y.value, 2)
935 with self.assertRaises(TypeError):
936 dumps(NEI.x)
937 with self.assertRaises(PicklingError):
938 dumps(NEI)
939
940 def test_tuple_subclass(self):
941 class SomeTuple(tuple, Enum):
942 first = (1, 'for the money')
943 second = (2, 'for the show')
944 third = (3, 'for the music')
945 self.assertIs(type(SomeTuple.first), SomeTuple)
946 self.assertIsInstance(SomeTuple.second, tuple)
947 self.assertEqual(SomeTuple.third, (3, 'for the music'))
948 globals()['SomeTuple'] = SomeTuple
949 self.assertIs(loads(dumps(SomeTuple.first)), SomeTuple.first)
950
951 def test_duplicate_values_give_unique_enum_items(self):
952 class AutoNumber(Enum):
953 first = ()
954 second = ()
955 third = ()
956 def __new__(cls):
957 value = len(cls.__members__) + 1
958 obj = object.__new__(cls)
Ethan Furman520ad572013-07-19 19:47:21 -0700959 obj._value_ = value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700960 return obj
961 def __int__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700962 return int(self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700963 self.assertEqual(
964 list(AutoNumber),
965 [AutoNumber.first, AutoNumber.second, AutoNumber.third],
966 )
967 self.assertEqual(int(AutoNumber.second), 2)
Ethan Furman2aa27322013-07-19 19:35:56 -0700968 self.assertEqual(AutoNumber.third.value, 3)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700969 self.assertIs(AutoNumber(1), AutoNumber.first)
970
971 def test_inherited_new_from_enhanced_enum(self):
972 class AutoNumber(Enum):
973 def __new__(cls):
974 value = len(cls.__members__) + 1
975 obj = object.__new__(cls)
Ethan Furman520ad572013-07-19 19:47:21 -0700976 obj._value_ = value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700977 return obj
978 def __int__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700979 return int(self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700980 class Color(AutoNumber):
981 red = ()
982 green = ()
983 blue = ()
984 self.assertEqual(list(Color), [Color.red, Color.green, Color.blue])
985 self.assertEqual(list(map(int, Color)), [1, 2, 3])
986
987 def test_inherited_new_from_mixed_enum(self):
988 class AutoNumber(IntEnum):
989 def __new__(cls):
990 value = len(cls.__members__) + 1
991 obj = int.__new__(cls, value)
Ethan Furman520ad572013-07-19 19:47:21 -0700992 obj._value_ = value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700993 return obj
994 class Color(AutoNumber):
995 red = ()
996 green = ()
997 blue = ()
998 self.assertEqual(list(Color), [Color.red, Color.green, Color.blue])
999 self.assertEqual(list(map(int, Color)), [1, 2, 3])
1000
1001 def test_ordered_mixin(self):
1002 class OrderedEnum(Enum):
1003 def __ge__(self, other):
1004 if self.__class__ is other.__class__:
Ethan Furman520ad572013-07-19 19:47:21 -07001005 return self._value_ >= other._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001006 return NotImplemented
1007 def __gt__(self, other):
1008 if self.__class__ is other.__class__:
Ethan Furman520ad572013-07-19 19:47:21 -07001009 return self._value_ > other._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001010 return NotImplemented
1011 def __le__(self, other):
1012 if self.__class__ is other.__class__:
Ethan Furman520ad572013-07-19 19:47:21 -07001013 return self._value_ <= other._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001014 return NotImplemented
1015 def __lt__(self, other):
1016 if self.__class__ is other.__class__:
Ethan Furman520ad572013-07-19 19:47:21 -07001017 return self._value_ < other._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001018 return NotImplemented
1019 class Grade(OrderedEnum):
1020 A = 5
1021 B = 4
1022 C = 3
1023 D = 2
1024 F = 1
1025 self.assertGreater(Grade.A, Grade.B)
1026 self.assertLessEqual(Grade.F, Grade.C)
1027 self.assertLess(Grade.D, Grade.A)
1028 self.assertGreaterEqual(Grade.B, Grade.B)
Ethan Furman520ad572013-07-19 19:47:21 -07001029
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001030 def test_extending2(self):
1031 class Shade(Enum):
1032 def shade(self):
1033 print(self.name)
1034 class Color(Shade):
1035 red = 1
1036 green = 2
1037 blue = 3
1038 with self.assertRaises(TypeError):
1039 class MoreColor(Color):
1040 cyan = 4
1041 magenta = 5
1042 yellow = 6
1043
1044 def test_extending3(self):
1045 class Shade(Enum):
1046 def shade(self):
1047 return self.name
1048 class Color(Shade):
1049 def hex(self):
1050 return '%s hexlified!' % self.value
1051 class MoreColor(Color):
1052 cyan = 4
1053 magenta = 5
1054 yellow = 6
1055 self.assertEqual(MoreColor.magenta.hex(), '5 hexlified!')
1056
1057
1058 def test_no_duplicates(self):
1059 class UniqueEnum(Enum):
1060 def __init__(self, *args):
1061 cls = self.__class__
1062 if any(self.value == e.value for e in cls):
1063 a = self.name
1064 e = cls(self.value).name
1065 raise ValueError(
1066 "aliases not allowed in UniqueEnum: %r --> %r"
1067 % (a, e)
1068 )
1069 class Color(UniqueEnum):
1070 red = 1
1071 green = 2
1072 blue = 3
1073 with self.assertRaises(ValueError):
1074 class Color(UniqueEnum):
1075 red = 1
1076 green = 2
1077 blue = 3
1078 grene = 2
1079
1080 def test_init(self):
1081 class Planet(Enum):
1082 MERCURY = (3.303e+23, 2.4397e6)
1083 VENUS = (4.869e+24, 6.0518e6)
1084 EARTH = (5.976e+24, 6.37814e6)
1085 MARS = (6.421e+23, 3.3972e6)
1086 JUPITER = (1.9e+27, 7.1492e7)
1087 SATURN = (5.688e+26, 6.0268e7)
1088 URANUS = (8.686e+25, 2.5559e7)
1089 NEPTUNE = (1.024e+26, 2.4746e7)
1090 def __init__(self, mass, radius):
1091 self.mass = mass # in kilograms
1092 self.radius = radius # in meters
1093 @property
1094 def surface_gravity(self):
1095 # universal gravitational constant (m3 kg-1 s-2)
1096 G = 6.67300E-11
1097 return G * self.mass / (self.radius * self.radius)
1098 self.assertEqual(round(Planet.EARTH.surface_gravity, 2), 9.80)
1099 self.assertEqual(Planet.EARTH.value, (5.976e+24, 6.37814e6))
1100
Ethan Furman2aa27322013-07-19 19:35:56 -07001101 def test_nonhash_value(self):
1102 class AutoNumberInAList(Enum):
1103 def __new__(cls):
1104 value = [len(cls.__members__) + 1]
1105 obj = object.__new__(cls)
Ethan Furman520ad572013-07-19 19:47:21 -07001106 obj._value_ = value
Ethan Furman2aa27322013-07-19 19:35:56 -07001107 return obj
1108 class ColorInAList(AutoNumberInAList):
1109 red = ()
1110 green = ()
1111 blue = ()
1112 self.assertEqual(list(ColorInAList), [ColorInAList.red, ColorInAList.green, ColorInAList.blue])
1113 self.assertEqual(ColorInAList.red.value, [1])
1114 self.assertEqual(ColorInAList([1]), ColorInAList.red)
1115
Ethan Furmanb41803e2013-07-25 13:50:45 -07001116 def test_conflicting_types_resolved_in_new(self):
1117 class LabelledIntEnum(int, Enum):
1118 def __new__(cls, *args):
1119 value, label = args
1120 obj = int.__new__(cls, value)
1121 obj.label = label
1122 obj._value_ = value
1123 return obj
1124
1125 class LabelledList(LabelledIntEnum):
1126 unprocessed = (1, "Unprocessed")
1127 payment_complete = (2, "Payment Complete")
1128
1129 self.assertEqual(list(LabelledList), [LabelledList.unprocessed, LabelledList.payment_complete])
1130 self.assertEqual(LabelledList.unprocessed, 1)
1131 self.assertEqual(LabelledList(1), LabelledList.unprocessed)
Ethan Furman2aa27322013-07-19 19:35:56 -07001132
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001133
Ethan Furmanf24bb352013-07-18 17:05:39 -07001134class TestUnique(unittest.TestCase):
1135
1136 def test_unique_clean(self):
1137 @unique
1138 class Clean(Enum):
1139 one = 1
1140 two = 'dos'
1141 tres = 4.0
1142 @unique
1143 class Cleaner(IntEnum):
1144 single = 1
1145 double = 2
1146 triple = 3
1147
1148 def test_unique_dirty(self):
1149 with self.assertRaisesRegex(ValueError, 'tres.*one'):
1150 @unique
1151 class Dirty(Enum):
1152 one = 1
1153 two = 'dos'
1154 tres = 1
1155 with self.assertRaisesRegex(
1156 ValueError,
1157 'double.*single.*turkey.*triple',
1158 ):
1159 @unique
1160 class Dirtier(IntEnum):
1161 single = 1
1162 double = 1
1163 triple = 3
1164 turkey = 3
1165
1166
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001167if __name__ == '__main__':
1168 unittest.main()