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