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