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