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