blob: 0f7b769b7afba69fd4cf07c4efbb90426446289b [file] [log] [blame]
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001import enum
Ethan Furman5875d742013-10-21 20:45:55 -07002import inspect
3import pydoc
Ethan Furman6b3d64a2013-06-14 16:55:46 -07004import unittest
5from collections import OrderedDict
Ethan Furman5875d742013-10-21 20:45:55 -07006from enum import Enum, IntEnum, EnumMeta, unique
7from io import StringIO
Ethan Furman2ddb39a2014-02-06 17:28:50 -08008from pickle import dumps, loads, PicklingError, HIGHEST_PROTOCOL
Ethan Furman6b3d64a2013-06-14 16:55:46 -07009
10# for pickle tests
11try:
12 class Stooges(Enum):
13 LARRY = 1
14 CURLY = 2
15 MOE = 3
16except Exception as exc:
17 Stooges = exc
18
19try:
20 class IntStooges(int, Enum):
21 LARRY = 1
22 CURLY = 2
23 MOE = 3
24except Exception as exc:
25 IntStooges = exc
26
27try:
28 class FloatStooges(float, Enum):
29 LARRY = 1.39
30 CURLY = 2.72
31 MOE = 3.142596
32except Exception as exc:
33 FloatStooges = exc
34
35# for pickle test and subclass tests
36try:
37 class StrEnum(str, Enum):
38 'accepts only string values'
39 class Name(StrEnum):
40 BDFL = 'Guido van Rossum'
41 FLUFL = 'Barry Warsaw'
42except Exception as exc:
43 Name = exc
44
45try:
46 Question = Enum('Question', 'who what when where why', module=__name__)
47except Exception as exc:
48 Question = exc
49
50try:
51 Answer = Enum('Answer', 'him this then there because')
52except Exception as exc:
53 Answer = exc
54
Ethan Furmanca1b7942014-02-08 11:36:27 -080055try:
56 Theory = Enum('Theory', 'rule law supposition', qualname='spanish_inquisition')
57except Exception as exc:
58 Theory = exc
59
Ethan Furman6b3d64a2013-06-14 16:55:46 -070060# for doctests
61try:
62 class Fruit(Enum):
63 tomato = 1
64 banana = 2
65 cherry = 3
66except Exception:
67 pass
68
Serhiy Storchakae50e7802015-03-31 16:56:49 +030069def test_pickle_dump_load(assertion, source, target=None):
Ethan Furman2ddb39a2014-02-06 17:28:50 -080070 if target is None:
71 target = source
Serhiy Storchakae50e7802015-03-31 16:56:49 +030072 for protocol in range(HIGHEST_PROTOCOL + 1):
Ethan Furman2ddb39a2014-02-06 17:28:50 -080073 assertion(loads(dumps(source, protocol=protocol)), target)
74
Serhiy Storchakae50e7802015-03-31 16:56:49 +030075def test_pickle_exception(assertion, exception, obj):
76 for protocol in range(HIGHEST_PROTOCOL + 1):
Ethan Furman2ddb39a2014-02-06 17:28:50 -080077 with assertion(exception):
78 dumps(obj, protocol=protocol)
Ethan Furman648f8602013-10-06 17:19:54 -070079
80class TestHelpers(unittest.TestCase):
81 # _is_descriptor, _is_sunder, _is_dunder
82
83 def test_is_descriptor(self):
84 class foo:
85 pass
86 for attr in ('__get__','__set__','__delete__'):
87 obj = foo()
88 self.assertFalse(enum._is_descriptor(obj))
89 setattr(obj, attr, 1)
90 self.assertTrue(enum._is_descriptor(obj))
91
92 def test_is_sunder(self):
93 for s in ('_a_', '_aa_'):
94 self.assertTrue(enum._is_sunder(s))
95
96 for s in ('a', 'a_', '_a', '__a', 'a__', '__a__', '_a__', '__a_', '_',
97 '__', '___', '____', '_____',):
98 self.assertFalse(enum._is_sunder(s))
99
100 def test_is_dunder(self):
101 for s in ('__a__', '__aa__'):
102 self.assertTrue(enum._is_dunder(s))
103 for s in ('a', 'a_', '_a', '__a', 'a__', '_a_', '_a__', '__a_', '_',
104 '__', '___', '____', '_____',):
105 self.assertFalse(enum._is_dunder(s))
106
107
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700108class TestEnum(unittest.TestCase):
Ethan Furmanca1b7942014-02-08 11:36:27 -0800109
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700110 def setUp(self):
111 class Season(Enum):
112 SPRING = 1
113 SUMMER = 2
114 AUTUMN = 3
115 WINTER = 4
116 self.Season = Season
117
Ethan Furmanec15a822013-08-31 19:17:41 -0700118 class Konstants(float, Enum):
119 E = 2.7182818
120 PI = 3.1415926
121 TAU = 2 * PI
122 self.Konstants = Konstants
123
124 class Grades(IntEnum):
125 A = 5
126 B = 4
127 C = 3
128 D = 2
129 F = 0
130 self.Grades = Grades
131
132 class Directional(str, Enum):
133 EAST = 'east'
134 WEST = 'west'
135 NORTH = 'north'
136 SOUTH = 'south'
137 self.Directional = Directional
138
139 from datetime import date
140 class Holiday(date, Enum):
141 NEW_YEAR = 2013, 1, 1
142 IDES_OF_MARCH = 2013, 3, 15
143 self.Holiday = Holiday
144
Ethan Furman388a3922013-08-12 06:51:41 -0700145 def test_dir_on_class(self):
146 Season = self.Season
147 self.assertEqual(
148 set(dir(Season)),
Ethan Furmanc850f342013-09-15 16:59:35 -0700149 set(['__class__', '__doc__', '__members__', '__module__',
Ethan Furman388a3922013-08-12 06:51:41 -0700150 'SPRING', 'SUMMER', 'AUTUMN', 'WINTER']),
151 )
152
153 def test_dir_on_item(self):
154 Season = self.Season
155 self.assertEqual(
156 set(dir(Season.WINTER)),
Ethan Furmanc850f342013-09-15 16:59:35 -0700157 set(['__class__', '__doc__', '__module__', 'name', 'value']),
Ethan Furman388a3922013-08-12 06:51:41 -0700158 )
159
Ethan Furmanc850f342013-09-15 16:59:35 -0700160 def test_dir_with_added_behavior(self):
161 class Test(Enum):
162 this = 'that'
163 these = 'those'
164 def wowser(self):
165 return ("Wowser! I'm %s!" % self.name)
166 self.assertEqual(
167 set(dir(Test)),
168 set(['__class__', '__doc__', '__members__', '__module__', 'this', 'these']),
169 )
170 self.assertEqual(
171 set(dir(Test.this)),
172 set(['__class__', '__doc__', '__module__', 'name', 'value', 'wowser']),
173 )
174
Ethan Furman0ae550b2014-10-14 08:58:32 -0700175 def test_dir_on_sub_with_behavior_on_super(self):
176 # see issue22506
177 class SuperEnum(Enum):
178 def invisible(self):
179 return "did you see me?"
180 class SubEnum(SuperEnum):
181 sample = 5
182 self.assertEqual(
183 set(dir(SubEnum.sample)),
184 set(['__class__', '__doc__', '__module__', 'name', 'value', 'invisible']),
185 )
186
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700187 def test_enum_in_enum_out(self):
188 Season = self.Season
189 self.assertIs(Season(Season.WINTER), Season.WINTER)
190
191 def test_enum_value(self):
192 Season = self.Season
193 self.assertEqual(Season.SPRING.value, 1)
194
195 def test_intenum_value(self):
196 self.assertEqual(IntStooges.CURLY.value, 2)
197
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700198 def test_enum(self):
199 Season = self.Season
200 lst = list(Season)
201 self.assertEqual(len(lst), len(Season))
202 self.assertEqual(len(Season), 4, Season)
203 self.assertEqual(
204 [Season.SPRING, Season.SUMMER, Season.AUTUMN, Season.WINTER], lst)
205
206 for i, season in enumerate('SPRING SUMMER AUTUMN WINTER'.split(), 1):
207 e = Season(i)
208 self.assertEqual(e, getattr(Season, season))
209 self.assertEqual(e.value, i)
210 self.assertNotEqual(e, i)
211 self.assertEqual(e.name, season)
212 self.assertIn(e, Season)
213 self.assertIs(type(e), Season)
214 self.assertIsInstance(e, Season)
215 self.assertEqual(str(e), 'Season.' + season)
216 self.assertEqual(
217 repr(e),
218 '<Season.{0}: {1}>'.format(season, i),
219 )
220
221 def test_value_name(self):
222 Season = self.Season
223 self.assertEqual(Season.SPRING.name, 'SPRING')
224 self.assertEqual(Season.SPRING.value, 1)
225 with self.assertRaises(AttributeError):
226 Season.SPRING.name = 'invierno'
227 with self.assertRaises(AttributeError):
228 Season.SPRING.value = 2
229
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700230 def test_changing_member(self):
231 Season = self.Season
232 with self.assertRaises(AttributeError):
233 Season.WINTER = 'really cold'
234
Ethan Furman64a99722013-09-22 16:18:19 -0700235 def test_attribute_deletion(self):
236 class Season(Enum):
237 SPRING = 1
238 SUMMER = 2
239 AUTUMN = 3
240 WINTER = 4
241
242 def spam(cls):
243 pass
244
245 self.assertTrue(hasattr(Season, 'spam'))
246 del Season.spam
247 self.assertFalse(hasattr(Season, 'spam'))
248
249 with self.assertRaises(AttributeError):
250 del Season.SPRING
251 with self.assertRaises(AttributeError):
252 del Season.DRY
253 with self.assertRaises(AttributeError):
254 del Season.SPRING.name
255
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700256 def test_invalid_names(self):
257 with self.assertRaises(ValueError):
258 class Wrong(Enum):
259 mro = 9
260 with self.assertRaises(ValueError):
261 class Wrong(Enum):
262 _create_= 11
263 with self.assertRaises(ValueError):
264 class Wrong(Enum):
265 _get_mixins_ = 9
266 with self.assertRaises(ValueError):
267 class Wrong(Enum):
268 _find_new_ = 1
269 with self.assertRaises(ValueError):
270 class Wrong(Enum):
271 _any_name_ = 9
272
Ethan Furman6db1fd52015-09-17 21:49:12 -0700273 def test_bool(self):
274 class Logic(Enum):
275 true = True
276 false = False
277 self.assertTrue(Logic.true)
278 self.assertFalse(Logic.false)
279
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700280 def test_contains(self):
281 Season = self.Season
282 self.assertIn(Season.AUTUMN, Season)
283 self.assertNotIn(3, Season)
284
285 val = Season(3)
286 self.assertIn(val, Season)
287
288 class OtherEnum(Enum):
289 one = 1; two = 2
290 self.assertNotIn(OtherEnum.two, Season)
291
292 def test_comparisons(self):
293 Season = self.Season
294 with self.assertRaises(TypeError):
295 Season.SPRING < Season.WINTER
296 with self.assertRaises(TypeError):
297 Season.SPRING > 4
298
299 self.assertNotEqual(Season.SPRING, 1)
300
301 class Part(Enum):
302 SPRING = 1
303 CLIP = 2
304 BARREL = 3
305
306 self.assertNotEqual(Season.SPRING, Part.SPRING)
307 with self.assertRaises(TypeError):
308 Season.SPRING < Part.CLIP
309
310 def test_enum_duplicates(self):
311 class Season(Enum):
312 SPRING = 1
313 SUMMER = 2
314 AUTUMN = FALL = 3
315 WINTER = 4
316 ANOTHER_SPRING = 1
317 lst = list(Season)
318 self.assertEqual(
319 lst,
320 [Season.SPRING, Season.SUMMER,
321 Season.AUTUMN, Season.WINTER,
322 ])
323 self.assertIs(Season.FALL, Season.AUTUMN)
324 self.assertEqual(Season.FALL.value, 3)
325 self.assertEqual(Season.AUTUMN.value, 3)
326 self.assertIs(Season(3), Season.AUTUMN)
327 self.assertIs(Season(1), Season.SPRING)
328 self.assertEqual(Season.FALL.name, 'AUTUMN')
329 self.assertEqual(
330 [k for k,v in Season.__members__.items() if v.name != k],
331 ['FALL', 'ANOTHER_SPRING'],
332 )
333
Ethan Furman101e0742013-09-15 12:34:36 -0700334 def test_duplicate_name(self):
335 with self.assertRaises(TypeError):
336 class Color(Enum):
337 red = 1
338 green = 2
339 blue = 3
340 red = 4
341
342 with self.assertRaises(TypeError):
343 class Color(Enum):
344 red = 1
345 green = 2
346 blue = 3
347 def red(self):
348 return 'red'
349
350 with self.assertRaises(TypeError):
351 class Color(Enum):
352 @property
353 def red(self):
354 return 'redder'
355 red = 1
356 green = 2
357 blue = 3
358
359
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700360 def test_enum_with_value_name(self):
361 class Huh(Enum):
362 name = 1
363 value = 2
364 self.assertEqual(
365 list(Huh),
366 [Huh.name, Huh.value],
367 )
368 self.assertIs(type(Huh.name), Huh)
369 self.assertEqual(Huh.name.name, 'name')
370 self.assertEqual(Huh.name.value, 1)
Ethan Furmanec15a822013-08-31 19:17:41 -0700371
372 def test_format_enum(self):
373 Season = self.Season
374 self.assertEqual('{}'.format(Season.SPRING),
375 '{}'.format(str(Season.SPRING)))
376 self.assertEqual( '{:}'.format(Season.SPRING),
377 '{:}'.format(str(Season.SPRING)))
378 self.assertEqual('{:20}'.format(Season.SPRING),
379 '{:20}'.format(str(Season.SPRING)))
380 self.assertEqual('{:^20}'.format(Season.SPRING),
381 '{:^20}'.format(str(Season.SPRING)))
382 self.assertEqual('{:>20}'.format(Season.SPRING),
383 '{:>20}'.format(str(Season.SPRING)))
384 self.assertEqual('{:<20}'.format(Season.SPRING),
385 '{:<20}'.format(str(Season.SPRING)))
386
387 def test_format_enum_custom(self):
388 class TestFloat(float, Enum):
389 one = 1.0
390 two = 2.0
391 def __format__(self, spec):
392 return 'TestFloat success!'
393 self.assertEqual('{}'.format(TestFloat.one), 'TestFloat success!')
394
395 def assertFormatIsValue(self, spec, member):
396 self.assertEqual(spec.format(member), spec.format(member.value))
397
398 def test_format_enum_date(self):
399 Holiday = self.Holiday
400 self.assertFormatIsValue('{}', Holiday.IDES_OF_MARCH)
401 self.assertFormatIsValue('{:}', Holiday.IDES_OF_MARCH)
402 self.assertFormatIsValue('{:20}', Holiday.IDES_OF_MARCH)
403 self.assertFormatIsValue('{:^20}', Holiday.IDES_OF_MARCH)
404 self.assertFormatIsValue('{:>20}', Holiday.IDES_OF_MARCH)
405 self.assertFormatIsValue('{:<20}', Holiday.IDES_OF_MARCH)
406 self.assertFormatIsValue('{:%Y %m}', Holiday.IDES_OF_MARCH)
407 self.assertFormatIsValue('{:%Y %m %M:00}', Holiday.IDES_OF_MARCH)
408
409 def test_format_enum_float(self):
410 Konstants = self.Konstants
411 self.assertFormatIsValue('{}', Konstants.TAU)
412 self.assertFormatIsValue('{:}', Konstants.TAU)
413 self.assertFormatIsValue('{:20}', Konstants.TAU)
414 self.assertFormatIsValue('{:^20}', Konstants.TAU)
415 self.assertFormatIsValue('{:>20}', Konstants.TAU)
416 self.assertFormatIsValue('{:<20}', Konstants.TAU)
417 self.assertFormatIsValue('{:n}', Konstants.TAU)
418 self.assertFormatIsValue('{:5.2}', Konstants.TAU)
419 self.assertFormatIsValue('{:f}', Konstants.TAU)
420
421 def test_format_enum_int(self):
422 Grades = self.Grades
423 self.assertFormatIsValue('{}', Grades.C)
424 self.assertFormatIsValue('{:}', Grades.C)
425 self.assertFormatIsValue('{:20}', Grades.C)
426 self.assertFormatIsValue('{:^20}', Grades.C)
427 self.assertFormatIsValue('{:>20}', Grades.C)
428 self.assertFormatIsValue('{:<20}', Grades.C)
429 self.assertFormatIsValue('{:+}', Grades.C)
430 self.assertFormatIsValue('{:08X}', Grades.C)
431 self.assertFormatIsValue('{:b}', Grades.C)
432
433 def test_format_enum_str(self):
434 Directional = self.Directional
435 self.assertFormatIsValue('{}', Directional.WEST)
436 self.assertFormatIsValue('{:}', Directional.WEST)
437 self.assertFormatIsValue('{:20}', Directional.WEST)
438 self.assertFormatIsValue('{:^20}', Directional.WEST)
439 self.assertFormatIsValue('{:>20}', Directional.WEST)
440 self.assertFormatIsValue('{:<20}', Directional.WEST)
441
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700442 def test_hash(self):
443 Season = self.Season
444 dates = {}
445 dates[Season.WINTER] = '1225'
446 dates[Season.SPRING] = '0315'
447 dates[Season.SUMMER] = '0704'
448 dates[Season.AUTUMN] = '1031'
449 self.assertEqual(dates[Season.AUTUMN], '1031')
450
451 def test_intenum_from_scratch(self):
452 class phy(int, Enum):
453 pi = 3
454 tau = 2 * pi
455 self.assertTrue(phy.pi < phy.tau)
456
457 def test_intenum_inherited(self):
458 class IntEnum(int, Enum):
459 pass
460 class phy(IntEnum):
461 pi = 3
462 tau = 2 * pi
463 self.assertTrue(phy.pi < phy.tau)
464
465 def test_floatenum_from_scratch(self):
466 class phy(float, Enum):
Ethan Furmanec15a822013-08-31 19:17:41 -0700467 pi = 3.1415926
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700468 tau = 2 * pi
469 self.assertTrue(phy.pi < phy.tau)
470
471 def test_floatenum_inherited(self):
472 class FloatEnum(float, Enum):
473 pass
474 class phy(FloatEnum):
Ethan Furmanec15a822013-08-31 19:17:41 -0700475 pi = 3.1415926
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700476 tau = 2 * pi
477 self.assertTrue(phy.pi < phy.tau)
478
479 def test_strenum_from_scratch(self):
480 class phy(str, Enum):
481 pi = 'Pi'
482 tau = 'Tau'
483 self.assertTrue(phy.pi < phy.tau)
484
485 def test_strenum_inherited(self):
486 class StrEnum(str, Enum):
487 pass
488 class phy(StrEnum):
489 pi = 'Pi'
490 tau = 'Tau'
491 self.assertTrue(phy.pi < phy.tau)
492
493
494 def test_intenum(self):
495 class WeekDay(IntEnum):
496 SUNDAY = 1
497 MONDAY = 2
498 TUESDAY = 3
499 WEDNESDAY = 4
500 THURSDAY = 5
501 FRIDAY = 6
502 SATURDAY = 7
503
504 self.assertEqual(['a', 'b', 'c'][WeekDay.MONDAY], 'c')
505 self.assertEqual([i for i in range(WeekDay.TUESDAY)], [0, 1, 2])
506
507 lst = list(WeekDay)
508 self.assertEqual(len(lst), len(WeekDay))
509 self.assertEqual(len(WeekDay), 7)
510 target = 'SUNDAY MONDAY TUESDAY WEDNESDAY THURSDAY FRIDAY SATURDAY'
511 target = target.split()
512 for i, weekday in enumerate(target, 1):
513 e = WeekDay(i)
514 self.assertEqual(e, i)
515 self.assertEqual(int(e), i)
516 self.assertEqual(e.name, weekday)
517 self.assertIn(e, WeekDay)
518 self.assertEqual(lst.index(e)+1, i)
519 self.assertTrue(0 < e < 8)
520 self.assertIs(type(e), WeekDay)
521 self.assertIsInstance(e, int)
522 self.assertIsInstance(e, Enum)
523
524 def test_intenum_duplicates(self):
525 class WeekDay(IntEnum):
526 SUNDAY = 1
527 MONDAY = 2
528 TUESDAY = TEUSDAY = 3
529 WEDNESDAY = 4
530 THURSDAY = 5
531 FRIDAY = 6
532 SATURDAY = 7
533 self.assertIs(WeekDay.TEUSDAY, WeekDay.TUESDAY)
534 self.assertEqual(WeekDay(3).name, 'TUESDAY')
535 self.assertEqual([k for k,v in WeekDay.__members__.items()
536 if v.name != k], ['TEUSDAY', ])
537
538 def test_pickle_enum(self):
539 if isinstance(Stooges, Exception):
540 raise Stooges
Ethan Furman2ddb39a2014-02-06 17:28:50 -0800541 test_pickle_dump_load(self.assertIs, Stooges.CURLY)
542 test_pickle_dump_load(self.assertIs, Stooges)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700543
544 def test_pickle_int(self):
545 if isinstance(IntStooges, Exception):
546 raise IntStooges
Ethan Furman2ddb39a2014-02-06 17:28:50 -0800547 test_pickle_dump_load(self.assertIs, IntStooges.CURLY)
548 test_pickle_dump_load(self.assertIs, IntStooges)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700549
550 def test_pickle_float(self):
551 if isinstance(FloatStooges, Exception):
552 raise FloatStooges
Ethan Furman2ddb39a2014-02-06 17:28:50 -0800553 test_pickle_dump_load(self.assertIs, FloatStooges.CURLY)
554 test_pickle_dump_load(self.assertIs, FloatStooges)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700555
556 def test_pickle_enum_function(self):
557 if isinstance(Answer, Exception):
558 raise Answer
Ethan Furman2ddb39a2014-02-06 17:28:50 -0800559 test_pickle_dump_load(self.assertIs, Answer.him)
560 test_pickle_dump_load(self.assertIs, Answer)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700561
562 def test_pickle_enum_function_with_module(self):
563 if isinstance(Question, Exception):
564 raise Question
Ethan Furman2ddb39a2014-02-06 17:28:50 -0800565 test_pickle_dump_load(self.assertIs, Question.who)
566 test_pickle_dump_load(self.assertIs, Question)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700567
Ethan Furmanca1b7942014-02-08 11:36:27 -0800568 def test_enum_function_with_qualname(self):
569 if isinstance(Theory, Exception):
570 raise Theory
571 self.assertEqual(Theory.__qualname__, 'spanish_inquisition')
572
573 def test_class_nested_enum_and_pickle_protocol_four(self):
574 # would normally just have this directly in the class namespace
575 class NestedEnum(Enum):
576 twigs = 'common'
577 shiny = 'rare'
578
579 self.__class__.NestedEnum = NestedEnum
580 self.NestedEnum.__qualname__ = '%s.NestedEnum' % self.__class__.__name__
Serhiy Storchakae50e7802015-03-31 16:56:49 +0300581 test_pickle_dump_load(self.assertIs, self.NestedEnum.twigs)
Ethan Furmanca1b7942014-02-08 11:36:27 -0800582
Ethan Furman24e837f2015-03-18 17:27:57 -0700583 def test_pickle_by_name(self):
584 class ReplaceGlobalInt(IntEnum):
585 ONE = 1
586 TWO = 2
587 ReplaceGlobalInt.__reduce_ex__ = enum._reduce_ex_by_name
588 for proto in range(HIGHEST_PROTOCOL):
589 self.assertEqual(ReplaceGlobalInt.TWO.__reduce_ex__(proto), 'TWO')
590
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700591 def test_exploding_pickle(self):
Ethan Furmanca1b7942014-02-08 11:36:27 -0800592 BadPickle = Enum(
593 'BadPickle', 'dill sweet bread-n-butter', module=__name__)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700594 globals()['BadPickle'] = BadPickle
Ethan Furmanca1b7942014-02-08 11:36:27 -0800595 # now break BadPickle to test exception raising
596 enum._make_class_unpicklable(BadPickle)
Ethan Furman2ddb39a2014-02-06 17:28:50 -0800597 test_pickle_exception(self.assertRaises, TypeError, BadPickle.dill)
598 test_pickle_exception(self.assertRaises, PicklingError, BadPickle)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700599
600 def test_string_enum(self):
601 class SkillLevel(str, Enum):
602 master = 'what is the sound of one hand clapping?'
603 journeyman = 'why did the chicken cross the road?'
604 apprentice = 'knock, knock!'
605 self.assertEqual(SkillLevel.apprentice, 'knock, knock!')
606
607 def test_getattr_getitem(self):
608 class Period(Enum):
609 morning = 1
610 noon = 2
611 evening = 3
612 night = 4
613 self.assertIs(Period(2), Period.noon)
614 self.assertIs(getattr(Period, 'night'), Period.night)
615 self.assertIs(Period['morning'], Period.morning)
616
617 def test_getattr_dunder(self):
618 Season = self.Season
619 self.assertTrue(getattr(Season, '__eq__'))
620
621 def test_iteration_order(self):
622 class Season(Enum):
623 SUMMER = 2
624 WINTER = 4
625 AUTUMN = 3
626 SPRING = 1
627 self.assertEqual(
628 list(Season),
629 [Season.SUMMER, Season.WINTER, Season.AUTUMN, Season.SPRING],
630 )
631
Ethan Furman2131a4a2013-09-14 18:11:24 -0700632 def test_reversed_iteration_order(self):
633 self.assertEqual(
634 list(reversed(self.Season)),
635 [self.Season.WINTER, self.Season.AUTUMN, self.Season.SUMMER,
636 self.Season.SPRING]
637 )
638
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700639 def test_programatic_function_string(self):
640 SummerMonth = Enum('SummerMonth', 'june july august')
641 lst = list(SummerMonth)
642 self.assertEqual(len(lst), len(SummerMonth))
643 self.assertEqual(len(SummerMonth), 3, SummerMonth)
644 self.assertEqual(
645 [SummerMonth.june, SummerMonth.july, SummerMonth.august],
646 lst,
647 )
648 for i, month in enumerate('june july august'.split(), 1):
649 e = SummerMonth(i)
650 self.assertEqual(int(e.value), i)
651 self.assertNotEqual(e, i)
652 self.assertEqual(e.name, month)
653 self.assertIn(e, SummerMonth)
654 self.assertIs(type(e), SummerMonth)
655
Ethan Furmand9925a12014-09-16 20:35:55 -0700656 def test_programatic_function_string_with_start(self):
657 SummerMonth = Enum('SummerMonth', 'june july august', start=10)
658 lst = list(SummerMonth)
659 self.assertEqual(len(lst), len(SummerMonth))
660 self.assertEqual(len(SummerMonth), 3, SummerMonth)
661 self.assertEqual(
662 [SummerMonth.june, SummerMonth.july, SummerMonth.august],
663 lst,
664 )
665 for i, month in enumerate('june july august'.split(), 10):
666 e = SummerMonth(i)
667 self.assertEqual(int(e.value), i)
668 self.assertNotEqual(e, i)
669 self.assertEqual(e.name, month)
670 self.assertIn(e, SummerMonth)
671 self.assertIs(type(e), SummerMonth)
672
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700673 def test_programatic_function_string_list(self):
674 SummerMonth = Enum('SummerMonth', ['june', 'july', 'august'])
675 lst = list(SummerMonth)
676 self.assertEqual(len(lst), len(SummerMonth))
677 self.assertEqual(len(SummerMonth), 3, SummerMonth)
678 self.assertEqual(
679 [SummerMonth.june, SummerMonth.july, SummerMonth.august],
680 lst,
681 )
682 for i, month in enumerate('june july august'.split(), 1):
683 e = SummerMonth(i)
684 self.assertEqual(int(e.value), i)
685 self.assertNotEqual(e, i)
686 self.assertEqual(e.name, month)
687 self.assertIn(e, SummerMonth)
688 self.assertIs(type(e), SummerMonth)
689
Ethan Furmand9925a12014-09-16 20:35:55 -0700690 def test_programatic_function_string_list_with_start(self):
691 SummerMonth = Enum('SummerMonth', ['june', 'july', 'august'], start=20)
692 lst = list(SummerMonth)
693 self.assertEqual(len(lst), len(SummerMonth))
694 self.assertEqual(len(SummerMonth), 3, SummerMonth)
695 self.assertEqual(
696 [SummerMonth.june, SummerMonth.july, SummerMonth.august],
697 lst,
698 )
699 for i, month in enumerate('june july august'.split(), 20):
700 e = SummerMonth(i)
701 self.assertEqual(int(e.value), i)
702 self.assertNotEqual(e, i)
703 self.assertEqual(e.name, month)
704 self.assertIn(e, SummerMonth)
705 self.assertIs(type(e), SummerMonth)
706
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700707 def test_programatic_function_iterable(self):
708 SummerMonth = Enum(
709 'SummerMonth',
710 (('june', 1), ('july', 2), ('august', 3))
711 )
712 lst = list(SummerMonth)
713 self.assertEqual(len(lst), len(SummerMonth))
714 self.assertEqual(len(SummerMonth), 3, SummerMonth)
715 self.assertEqual(
716 [SummerMonth.june, SummerMonth.july, SummerMonth.august],
717 lst,
718 )
719 for i, month in enumerate('june july august'.split(), 1):
720 e = SummerMonth(i)
721 self.assertEqual(int(e.value), i)
722 self.assertNotEqual(e, i)
723 self.assertEqual(e.name, month)
724 self.assertIn(e, SummerMonth)
725 self.assertIs(type(e), SummerMonth)
726
727 def test_programatic_function_from_dict(self):
728 SummerMonth = Enum(
729 'SummerMonth',
730 OrderedDict((('june', 1), ('july', 2), ('august', 3)))
731 )
732 lst = list(SummerMonth)
733 self.assertEqual(len(lst), len(SummerMonth))
734 self.assertEqual(len(SummerMonth), 3, SummerMonth)
735 self.assertEqual(
736 [SummerMonth.june, SummerMonth.july, SummerMonth.august],
737 lst,
738 )
739 for i, month in enumerate('june july august'.split(), 1):
740 e = SummerMonth(i)
741 self.assertEqual(int(e.value), i)
742 self.assertNotEqual(e, i)
743 self.assertEqual(e.name, month)
744 self.assertIn(e, SummerMonth)
745 self.assertIs(type(e), SummerMonth)
746
747 def test_programatic_function_type(self):
748 SummerMonth = Enum('SummerMonth', 'june july august', type=int)
749 lst = list(SummerMonth)
750 self.assertEqual(len(lst), len(SummerMonth))
751 self.assertEqual(len(SummerMonth), 3, SummerMonth)
752 self.assertEqual(
753 [SummerMonth.june, SummerMonth.july, SummerMonth.august],
754 lst,
755 )
756 for i, month in enumerate('june july august'.split(), 1):
757 e = SummerMonth(i)
758 self.assertEqual(e, i)
759 self.assertEqual(e.name, month)
760 self.assertIn(e, SummerMonth)
761 self.assertIs(type(e), SummerMonth)
762
Ethan Furmand9925a12014-09-16 20:35:55 -0700763 def test_programatic_function_type_with_start(self):
764 SummerMonth = Enum('SummerMonth', 'june july august', type=int, start=30)
765 lst = list(SummerMonth)
766 self.assertEqual(len(lst), len(SummerMonth))
767 self.assertEqual(len(SummerMonth), 3, SummerMonth)
768 self.assertEqual(
769 [SummerMonth.june, SummerMonth.july, SummerMonth.august],
770 lst,
771 )
772 for i, month in enumerate('june july august'.split(), 30):
773 e = SummerMonth(i)
774 self.assertEqual(e, i)
775 self.assertEqual(e.name, month)
776 self.assertIn(e, SummerMonth)
777 self.assertIs(type(e), SummerMonth)
778
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700779 def test_programatic_function_type_from_subclass(self):
780 SummerMonth = IntEnum('SummerMonth', 'june july august')
781 lst = list(SummerMonth)
782 self.assertEqual(len(lst), len(SummerMonth))
783 self.assertEqual(len(SummerMonth), 3, SummerMonth)
784 self.assertEqual(
785 [SummerMonth.june, SummerMonth.july, SummerMonth.august],
786 lst,
787 )
788 for i, month in enumerate('june july august'.split(), 1):
789 e = SummerMonth(i)
790 self.assertEqual(e, i)
791 self.assertEqual(e.name, month)
792 self.assertIn(e, SummerMonth)
793 self.assertIs(type(e), SummerMonth)
794
Ethan Furmand9925a12014-09-16 20:35:55 -0700795 def test_programatic_function_type_from_subclass_with_start(self):
796 SummerMonth = IntEnum('SummerMonth', 'june july august', start=40)
797 lst = list(SummerMonth)
798 self.assertEqual(len(lst), len(SummerMonth))
799 self.assertEqual(len(SummerMonth), 3, SummerMonth)
800 self.assertEqual(
801 [SummerMonth.june, SummerMonth.july, SummerMonth.august],
802 lst,
803 )
804 for i, month in enumerate('june july august'.split(), 40):
805 e = SummerMonth(i)
806 self.assertEqual(e, i)
807 self.assertEqual(e.name, month)
808 self.assertIn(e, SummerMonth)
809 self.assertIs(type(e), SummerMonth)
810
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700811 def test_subclassing(self):
812 if isinstance(Name, Exception):
813 raise Name
814 self.assertEqual(Name.BDFL, 'Guido van Rossum')
815 self.assertTrue(Name.BDFL, Name('Guido van Rossum'))
816 self.assertIs(Name.BDFL, getattr(Name, 'BDFL'))
Ethan Furman2ddb39a2014-02-06 17:28:50 -0800817 test_pickle_dump_load(self.assertIs, Name.BDFL)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700818
819 def test_extending(self):
820 class Color(Enum):
821 red = 1
822 green = 2
823 blue = 3
824 with self.assertRaises(TypeError):
825 class MoreColor(Color):
826 cyan = 4
827 magenta = 5
828 yellow = 6
829
830 def test_exclude_methods(self):
831 class whatever(Enum):
832 this = 'that'
833 these = 'those'
834 def really(self):
835 return 'no, not %s' % self.value
836 self.assertIsNot(type(whatever.really), whatever)
837 self.assertEqual(whatever.this.really(), 'no, not that')
838
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700839 def test_wrong_inheritance_order(self):
840 with self.assertRaises(TypeError):
841 class Wrong(Enum, str):
842 NotHere = 'error before this point'
843
844 def test_intenum_transitivity(self):
845 class number(IntEnum):
846 one = 1
847 two = 2
848 three = 3
849 class numero(IntEnum):
850 uno = 1
851 dos = 2
852 tres = 3
853 self.assertEqual(number.one, numero.uno)
854 self.assertEqual(number.two, numero.dos)
855 self.assertEqual(number.three, numero.tres)
856
857 def test_wrong_enum_in_call(self):
858 class Monochrome(Enum):
859 black = 0
860 white = 1
861 class Gender(Enum):
862 male = 0
863 female = 1
864 self.assertRaises(ValueError, Monochrome, Gender.male)
865
866 def test_wrong_enum_in_mixed_call(self):
867 class Monochrome(IntEnum):
868 black = 0
869 white = 1
870 class Gender(Enum):
871 male = 0
872 female = 1
873 self.assertRaises(ValueError, Monochrome, Gender.male)
874
875 def test_mixed_enum_in_call_1(self):
876 class Monochrome(IntEnum):
877 black = 0
878 white = 1
879 class Gender(IntEnum):
880 male = 0
881 female = 1
882 self.assertIs(Monochrome(Gender.female), Monochrome.white)
883
884 def test_mixed_enum_in_call_2(self):
885 class Monochrome(Enum):
886 black = 0
887 white = 1
888 class Gender(IntEnum):
889 male = 0
890 female = 1
891 self.assertIs(Monochrome(Gender.male), Monochrome.black)
892
893 def test_flufl_enum(self):
894 class Fluflnum(Enum):
895 def __int__(self):
896 return int(self.value)
897 class MailManOptions(Fluflnum):
898 option1 = 1
899 option2 = 2
900 option3 = 3
901 self.assertEqual(int(MailManOptions.option1), 1)
902
Ethan Furman5e5a8232013-08-04 08:42:23 -0700903 def test_introspection(self):
904 class Number(IntEnum):
905 one = 100
906 two = 200
907 self.assertIs(Number.one._member_type_, int)
908 self.assertIs(Number._member_type_, int)
909 class String(str, Enum):
910 yarn = 'soft'
911 rope = 'rough'
912 wire = 'hard'
913 self.assertIs(String.yarn._member_type_, str)
914 self.assertIs(String._member_type_, str)
915 class Plain(Enum):
916 vanilla = 'white'
917 one = 1
918 self.assertIs(Plain.vanilla._member_type_, object)
919 self.assertIs(Plain._member_type_, object)
920
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700921 def test_no_such_enum_member(self):
922 class Color(Enum):
923 red = 1
924 green = 2
925 blue = 3
926 with self.assertRaises(ValueError):
927 Color(4)
928 with self.assertRaises(KeyError):
929 Color['chartreuse']
930
931 def test_new_repr(self):
932 class Color(Enum):
933 red = 1
934 green = 2
935 blue = 3
936 def __repr__(self):
937 return "don't you just love shades of %s?" % self.name
938 self.assertEqual(
939 repr(Color.blue),
940 "don't you just love shades of blue?",
941 )
942
943 def test_inherited_repr(self):
944 class MyEnum(Enum):
945 def __repr__(self):
946 return "My name is %s." % self.name
947 class MyIntEnum(int, MyEnum):
948 this = 1
949 that = 2
950 theother = 3
951 self.assertEqual(repr(MyIntEnum.that), "My name is that.")
952
953 def test_multiple_mixin_mro(self):
954 class auto_enum(type(Enum)):
955 def __new__(metacls, cls, bases, classdict):
956 temp = type(classdict)()
957 names = set(classdict._member_names)
958 i = 0
959 for k in classdict._member_names:
960 v = classdict[k]
961 if v is Ellipsis:
962 v = i
963 else:
964 i = v
965 i += 1
966 temp[k] = v
967 for k, v in classdict.items():
968 if k not in names:
969 temp[k] = v
970 return super(auto_enum, metacls).__new__(
971 metacls, cls, bases, temp)
972
973 class AutoNumberedEnum(Enum, metaclass=auto_enum):
974 pass
975
976 class AutoIntEnum(IntEnum, metaclass=auto_enum):
977 pass
978
979 class TestAutoNumber(AutoNumberedEnum):
980 a = ...
981 b = 3
982 c = ...
983
984 class TestAutoInt(AutoIntEnum):
985 a = ...
986 b = 3
987 c = ...
988
989 def test_subclasses_with_getnewargs(self):
990 class NamedInt(int):
Ethan Furman2ddb39a2014-02-06 17:28:50 -0800991 __qualname__ = 'NamedInt' # needed for pickle protocol 4
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700992 def __new__(cls, *args):
993 _args = args
994 name, *args = args
995 if len(args) == 0:
996 raise TypeError("name and value must be specified")
997 self = int.__new__(cls, *args)
998 self._intname = name
999 self._args = _args
1000 return self
1001 def __getnewargs__(self):
1002 return self._args
1003 @property
1004 def __name__(self):
1005 return self._intname
1006 def __repr__(self):
1007 # repr() is updated to include the name and type info
1008 return "{}({!r}, {})".format(type(self).__name__,
1009 self.__name__,
1010 int.__repr__(self))
1011 def __str__(self):
1012 # str() is unchanged, even if it relies on the repr() fallback
1013 base = int
1014 base_str = base.__str__
1015 if base_str.__objclass__ is object:
1016 return base.__repr__(self)
1017 return base_str(self)
1018 # for simplicity, we only define one operator that
1019 # propagates expressions
1020 def __add__(self, other):
1021 temp = int(self) + int( other)
1022 if isinstance(self, NamedInt) and isinstance(other, NamedInt):
1023 return NamedInt(
1024 '({0} + {1})'.format(self.__name__, other.__name__),
1025 temp )
1026 else:
1027 return temp
1028
1029 class NEI(NamedInt, Enum):
Ethan Furman2ddb39a2014-02-06 17:28:50 -08001030 __qualname__ = 'NEI' # needed for pickle protocol 4
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001031 x = ('the-x', 1)
1032 y = ('the-y', 2)
1033
Ethan Furman2aa27322013-07-19 19:35:56 -07001034
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001035 self.assertIs(NEI.__new__, Enum.__new__)
1036 self.assertEqual(repr(NEI.x + NEI.y), "NamedInt('(the-x + the-y)', 3)")
1037 globals()['NamedInt'] = NamedInt
1038 globals()['NEI'] = NEI
1039 NI5 = NamedInt('test', 5)
1040 self.assertEqual(NI5, 5)
Ethan Furman2ddb39a2014-02-06 17:28:50 -08001041 test_pickle_dump_load(self.assertEqual, NI5, 5)
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001042 self.assertEqual(NEI.y.value, 2)
Ethan Furman2ddb39a2014-02-06 17:28:50 -08001043 test_pickle_dump_load(self.assertIs, NEI.y)
Ethan Furmandc870522014-02-18 12:37:12 -08001044 test_pickle_dump_load(self.assertIs, NEI)
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001045
Ethan Furmanca1b7942014-02-08 11:36:27 -08001046 def test_subclasses_with_getnewargs_ex(self):
1047 class NamedInt(int):
1048 __qualname__ = 'NamedInt' # needed for pickle protocol 4
1049 def __new__(cls, *args):
1050 _args = args
1051 name, *args = args
1052 if len(args) == 0:
1053 raise TypeError("name and value must be specified")
1054 self = int.__new__(cls, *args)
1055 self._intname = name
1056 self._args = _args
1057 return self
1058 def __getnewargs_ex__(self):
1059 return self._args, {}
1060 @property
1061 def __name__(self):
1062 return self._intname
1063 def __repr__(self):
1064 # repr() is updated to include the name and type info
1065 return "{}({!r}, {})".format(type(self).__name__,
1066 self.__name__,
1067 int.__repr__(self))
1068 def __str__(self):
1069 # str() is unchanged, even if it relies on the repr() fallback
1070 base = int
1071 base_str = base.__str__
1072 if base_str.__objclass__ is object:
1073 return base.__repr__(self)
1074 return base_str(self)
1075 # for simplicity, we only define one operator that
1076 # propagates expressions
1077 def __add__(self, other):
1078 temp = int(self) + int( other)
1079 if isinstance(self, NamedInt) and isinstance(other, NamedInt):
1080 return NamedInt(
1081 '({0} + {1})'.format(self.__name__, other.__name__),
1082 temp )
1083 else:
1084 return temp
1085
1086 class NEI(NamedInt, Enum):
1087 __qualname__ = 'NEI' # needed for pickle protocol 4
1088 x = ('the-x', 1)
1089 y = ('the-y', 2)
1090
1091
1092 self.assertIs(NEI.__new__, Enum.__new__)
1093 self.assertEqual(repr(NEI.x + NEI.y), "NamedInt('(the-x + the-y)', 3)")
1094 globals()['NamedInt'] = NamedInt
1095 globals()['NEI'] = NEI
1096 NI5 = NamedInt('test', 5)
1097 self.assertEqual(NI5, 5)
Serhiy Storchakae50e7802015-03-31 16:56:49 +03001098 test_pickle_dump_load(self.assertEqual, NI5, 5)
Ethan Furmanca1b7942014-02-08 11:36:27 -08001099 self.assertEqual(NEI.y.value, 2)
Serhiy Storchakae50e7802015-03-31 16:56:49 +03001100 test_pickle_dump_load(self.assertIs, NEI.y)
Ethan Furmandc870522014-02-18 12:37:12 -08001101 test_pickle_dump_load(self.assertIs, NEI)
Ethan Furmanca1b7942014-02-08 11:36:27 -08001102
1103 def test_subclasses_with_reduce(self):
1104 class NamedInt(int):
1105 __qualname__ = 'NamedInt' # needed for pickle protocol 4
1106 def __new__(cls, *args):
1107 _args = args
1108 name, *args = args
1109 if len(args) == 0:
1110 raise TypeError("name and value must be specified")
1111 self = int.__new__(cls, *args)
1112 self._intname = name
1113 self._args = _args
1114 return self
1115 def __reduce__(self):
1116 return self.__class__, self._args
1117 @property
1118 def __name__(self):
1119 return self._intname
1120 def __repr__(self):
1121 # repr() is updated to include the name and type info
1122 return "{}({!r}, {})".format(type(self).__name__,
1123 self.__name__,
1124 int.__repr__(self))
1125 def __str__(self):
1126 # str() is unchanged, even if it relies on the repr() fallback
1127 base = int
1128 base_str = base.__str__
1129 if base_str.__objclass__ is object:
1130 return base.__repr__(self)
1131 return base_str(self)
1132 # for simplicity, we only define one operator that
1133 # propagates expressions
1134 def __add__(self, other):
1135 temp = int(self) + int( other)
1136 if isinstance(self, NamedInt) and isinstance(other, NamedInt):
1137 return NamedInt(
1138 '({0} + {1})'.format(self.__name__, other.__name__),
1139 temp )
1140 else:
1141 return temp
1142
1143 class NEI(NamedInt, Enum):
1144 __qualname__ = 'NEI' # needed for pickle protocol 4
1145 x = ('the-x', 1)
1146 y = ('the-y', 2)
1147
1148
1149 self.assertIs(NEI.__new__, Enum.__new__)
1150 self.assertEqual(repr(NEI.x + NEI.y), "NamedInt('(the-x + the-y)', 3)")
1151 globals()['NamedInt'] = NamedInt
1152 globals()['NEI'] = NEI
1153 NI5 = NamedInt('test', 5)
1154 self.assertEqual(NI5, 5)
1155 test_pickle_dump_load(self.assertEqual, NI5, 5)
1156 self.assertEqual(NEI.y.value, 2)
1157 test_pickle_dump_load(self.assertIs, NEI.y)
Ethan Furmandc870522014-02-18 12:37:12 -08001158 test_pickle_dump_load(self.assertIs, NEI)
Ethan Furmanca1b7942014-02-08 11:36:27 -08001159
1160 def test_subclasses_with_reduce_ex(self):
1161 class NamedInt(int):
1162 __qualname__ = 'NamedInt' # needed for pickle protocol 4
1163 def __new__(cls, *args):
1164 _args = args
1165 name, *args = args
1166 if len(args) == 0:
1167 raise TypeError("name and value must be specified")
1168 self = int.__new__(cls, *args)
1169 self._intname = name
1170 self._args = _args
1171 return self
1172 def __reduce_ex__(self, proto):
1173 return self.__class__, self._args
1174 @property
1175 def __name__(self):
1176 return self._intname
1177 def __repr__(self):
1178 # repr() is updated to include the name and type info
1179 return "{}({!r}, {})".format(type(self).__name__,
1180 self.__name__,
1181 int.__repr__(self))
1182 def __str__(self):
1183 # str() is unchanged, even if it relies on the repr() fallback
1184 base = int
1185 base_str = base.__str__
1186 if base_str.__objclass__ is object:
1187 return base.__repr__(self)
1188 return base_str(self)
1189 # for simplicity, we only define one operator that
1190 # propagates expressions
1191 def __add__(self, other):
1192 temp = int(self) + int( other)
1193 if isinstance(self, NamedInt) and isinstance(other, NamedInt):
1194 return NamedInt(
1195 '({0} + {1})'.format(self.__name__, other.__name__),
1196 temp )
1197 else:
1198 return temp
1199
1200 class NEI(NamedInt, Enum):
1201 __qualname__ = 'NEI' # needed for pickle protocol 4
1202 x = ('the-x', 1)
1203 y = ('the-y', 2)
1204
1205
1206 self.assertIs(NEI.__new__, Enum.__new__)
1207 self.assertEqual(repr(NEI.x + NEI.y), "NamedInt('(the-x + the-y)', 3)")
1208 globals()['NamedInt'] = NamedInt
1209 globals()['NEI'] = NEI
1210 NI5 = NamedInt('test', 5)
1211 self.assertEqual(NI5, 5)
1212 test_pickle_dump_load(self.assertEqual, NI5, 5)
1213 self.assertEqual(NEI.y.value, 2)
1214 test_pickle_dump_load(self.assertIs, NEI.y)
Ethan Furmandc870522014-02-18 12:37:12 -08001215 test_pickle_dump_load(self.assertIs, NEI)
Ethan Furmanca1b7942014-02-08 11:36:27 -08001216
Ethan Furmandc870522014-02-18 12:37:12 -08001217 def test_subclasses_without_direct_pickle_support(self):
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001218 class NamedInt(int):
Ethan Furman2ddb39a2014-02-06 17:28:50 -08001219 __qualname__ = 'NamedInt'
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001220 def __new__(cls, *args):
1221 _args = args
1222 name, *args = args
1223 if len(args) == 0:
1224 raise TypeError("name and value must be specified")
1225 self = int.__new__(cls, *args)
1226 self._intname = name
1227 self._args = _args
1228 return self
1229 @property
1230 def __name__(self):
1231 return self._intname
1232 def __repr__(self):
1233 # repr() is updated to include the name and type info
1234 return "{}({!r}, {})".format(type(self).__name__,
1235 self.__name__,
1236 int.__repr__(self))
1237 def __str__(self):
1238 # str() is unchanged, even if it relies on the repr() fallback
1239 base = int
1240 base_str = base.__str__
1241 if base_str.__objclass__ is object:
1242 return base.__repr__(self)
1243 return base_str(self)
1244 # for simplicity, we only define one operator that
1245 # propagates expressions
1246 def __add__(self, other):
1247 temp = int(self) + int( other)
1248 if isinstance(self, NamedInt) and isinstance(other, NamedInt):
1249 return NamedInt(
1250 '({0} + {1})'.format(self.__name__, other.__name__),
1251 temp )
1252 else:
1253 return temp
1254
1255 class NEI(NamedInt, Enum):
Ethan Furman2ddb39a2014-02-06 17:28:50 -08001256 __qualname__ = 'NEI'
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001257 x = ('the-x', 1)
1258 y = ('the-y', 2)
1259
1260 self.assertIs(NEI.__new__, Enum.__new__)
1261 self.assertEqual(repr(NEI.x + NEI.y), "NamedInt('(the-x + the-y)', 3)")
1262 globals()['NamedInt'] = NamedInt
1263 globals()['NEI'] = NEI
1264 NI5 = NamedInt('test', 5)
1265 self.assertEqual(NI5, 5)
1266 self.assertEqual(NEI.y.value, 2)
Ethan Furman2ddb39a2014-02-06 17:28:50 -08001267 test_pickle_exception(self.assertRaises, TypeError, NEI.x)
1268 test_pickle_exception(self.assertRaises, PicklingError, NEI)
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001269
Ethan Furmandc870522014-02-18 12:37:12 -08001270 def test_subclasses_without_direct_pickle_support_using_name(self):
1271 class NamedInt(int):
1272 __qualname__ = 'NamedInt'
1273 def __new__(cls, *args):
1274 _args = args
1275 name, *args = args
1276 if len(args) == 0:
1277 raise TypeError("name and value must be specified")
1278 self = int.__new__(cls, *args)
1279 self._intname = name
1280 self._args = _args
1281 return self
1282 @property
1283 def __name__(self):
1284 return self._intname
1285 def __repr__(self):
1286 # repr() is updated to include the name and type info
1287 return "{}({!r}, {})".format(type(self).__name__,
1288 self.__name__,
1289 int.__repr__(self))
1290 def __str__(self):
1291 # str() is unchanged, even if it relies on the repr() fallback
1292 base = int
1293 base_str = base.__str__
1294 if base_str.__objclass__ is object:
1295 return base.__repr__(self)
1296 return base_str(self)
1297 # for simplicity, we only define one operator that
1298 # propagates expressions
1299 def __add__(self, other):
1300 temp = int(self) + int( other)
1301 if isinstance(self, NamedInt) and isinstance(other, NamedInt):
1302 return NamedInt(
1303 '({0} + {1})'.format(self.__name__, other.__name__),
1304 temp )
1305 else:
1306 return temp
1307
1308 class NEI(NamedInt, Enum):
1309 __qualname__ = 'NEI'
1310 x = ('the-x', 1)
1311 y = ('the-y', 2)
1312 def __reduce_ex__(self, proto):
1313 return getattr, (self.__class__, self._name_)
1314
1315 self.assertIs(NEI.__new__, Enum.__new__)
1316 self.assertEqual(repr(NEI.x + NEI.y), "NamedInt('(the-x + the-y)', 3)")
1317 globals()['NamedInt'] = NamedInt
1318 globals()['NEI'] = NEI
1319 NI5 = NamedInt('test', 5)
1320 self.assertEqual(NI5, 5)
1321 self.assertEqual(NEI.y.value, 2)
1322 test_pickle_dump_load(self.assertIs, NEI.y)
1323 test_pickle_dump_load(self.assertIs, NEI)
1324
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001325 def test_tuple_subclass(self):
1326 class SomeTuple(tuple, Enum):
Ethan Furman2ddb39a2014-02-06 17:28:50 -08001327 __qualname__ = 'SomeTuple' # needed for pickle protocol 4
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001328 first = (1, 'for the money')
1329 second = (2, 'for the show')
1330 third = (3, 'for the music')
1331 self.assertIs(type(SomeTuple.first), SomeTuple)
1332 self.assertIsInstance(SomeTuple.second, tuple)
1333 self.assertEqual(SomeTuple.third, (3, 'for the music'))
1334 globals()['SomeTuple'] = SomeTuple
Ethan Furman2ddb39a2014-02-06 17:28:50 -08001335 test_pickle_dump_load(self.assertIs, SomeTuple.first)
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001336
1337 def test_duplicate_values_give_unique_enum_items(self):
1338 class AutoNumber(Enum):
1339 first = ()
1340 second = ()
1341 third = ()
1342 def __new__(cls):
1343 value = len(cls.__members__) + 1
1344 obj = object.__new__(cls)
Ethan Furman520ad572013-07-19 19:47:21 -07001345 obj._value_ = value
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001346 return obj
1347 def __int__(self):
Ethan Furman520ad572013-07-19 19:47:21 -07001348 return int(self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001349 self.assertEqual(
1350 list(AutoNumber),
1351 [AutoNumber.first, AutoNumber.second, AutoNumber.third],
1352 )
1353 self.assertEqual(int(AutoNumber.second), 2)
Ethan Furman2aa27322013-07-19 19:35:56 -07001354 self.assertEqual(AutoNumber.third.value, 3)
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001355 self.assertIs(AutoNumber(1), AutoNumber.first)
1356
1357 def test_inherited_new_from_enhanced_enum(self):
1358 class AutoNumber(Enum):
1359 def __new__(cls):
1360 value = len(cls.__members__) + 1
1361 obj = object.__new__(cls)
Ethan Furman520ad572013-07-19 19:47:21 -07001362 obj._value_ = value
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001363 return obj
1364 def __int__(self):
Ethan Furman520ad572013-07-19 19:47:21 -07001365 return int(self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001366 class Color(AutoNumber):
1367 red = ()
1368 green = ()
1369 blue = ()
1370 self.assertEqual(list(Color), [Color.red, Color.green, Color.blue])
1371 self.assertEqual(list(map(int, Color)), [1, 2, 3])
1372
1373 def test_inherited_new_from_mixed_enum(self):
1374 class AutoNumber(IntEnum):
1375 def __new__(cls):
1376 value = len(cls.__members__) + 1
1377 obj = int.__new__(cls, value)
Ethan Furman520ad572013-07-19 19:47:21 -07001378 obj._value_ = value
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001379 return obj
1380 class Color(AutoNumber):
1381 red = ()
1382 green = ()
1383 blue = ()
1384 self.assertEqual(list(Color), [Color.red, Color.green, Color.blue])
1385 self.assertEqual(list(map(int, Color)), [1, 2, 3])
1386
Ethan Furmanbe3c2fe2013-11-13 14:25:45 -08001387 def test_equality(self):
1388 class AlwaysEqual:
1389 def __eq__(self, other):
1390 return True
1391 class OrdinaryEnum(Enum):
1392 a = 1
1393 self.assertEqual(AlwaysEqual(), OrdinaryEnum.a)
1394 self.assertEqual(OrdinaryEnum.a, AlwaysEqual())
1395
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001396 def test_ordered_mixin(self):
1397 class OrderedEnum(Enum):
1398 def __ge__(self, other):
1399 if self.__class__ is other.__class__:
Ethan Furman520ad572013-07-19 19:47:21 -07001400 return self._value_ >= other._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001401 return NotImplemented
1402 def __gt__(self, other):
1403 if self.__class__ is other.__class__:
Ethan Furman520ad572013-07-19 19:47:21 -07001404 return self._value_ > other._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001405 return NotImplemented
1406 def __le__(self, other):
1407 if self.__class__ is other.__class__:
Ethan Furman520ad572013-07-19 19:47:21 -07001408 return self._value_ <= other._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001409 return NotImplemented
1410 def __lt__(self, other):
1411 if self.__class__ is other.__class__:
Ethan Furman520ad572013-07-19 19:47:21 -07001412 return self._value_ < other._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001413 return NotImplemented
1414 class Grade(OrderedEnum):
1415 A = 5
1416 B = 4
1417 C = 3
1418 D = 2
1419 F = 1
1420 self.assertGreater(Grade.A, Grade.B)
1421 self.assertLessEqual(Grade.F, Grade.C)
1422 self.assertLess(Grade.D, Grade.A)
1423 self.assertGreaterEqual(Grade.B, Grade.B)
Ethan Furmanbe3c2fe2013-11-13 14:25:45 -08001424 self.assertEqual(Grade.B, Grade.B)
1425 self.assertNotEqual(Grade.C, Grade.D)
Ethan Furman520ad572013-07-19 19:47:21 -07001426
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001427 def test_extending2(self):
1428 class Shade(Enum):
1429 def shade(self):
1430 print(self.name)
1431 class Color(Shade):
1432 red = 1
1433 green = 2
1434 blue = 3
1435 with self.assertRaises(TypeError):
1436 class MoreColor(Color):
1437 cyan = 4
1438 magenta = 5
1439 yellow = 6
1440
1441 def test_extending3(self):
1442 class Shade(Enum):
1443 def shade(self):
1444 return self.name
1445 class Color(Shade):
1446 def hex(self):
1447 return '%s hexlified!' % self.value
1448 class MoreColor(Color):
1449 cyan = 4
1450 magenta = 5
1451 yellow = 6
1452 self.assertEqual(MoreColor.magenta.hex(), '5 hexlified!')
1453
1454
1455 def test_no_duplicates(self):
1456 class UniqueEnum(Enum):
1457 def __init__(self, *args):
1458 cls = self.__class__
1459 if any(self.value == e.value for e in cls):
1460 a = self.name
1461 e = cls(self.value).name
1462 raise ValueError(
1463 "aliases not allowed in UniqueEnum: %r --> %r"
1464 % (a, e)
1465 )
1466 class Color(UniqueEnum):
1467 red = 1
1468 green = 2
1469 blue = 3
1470 with self.assertRaises(ValueError):
1471 class Color(UniqueEnum):
1472 red = 1
1473 green = 2
1474 blue = 3
1475 grene = 2
1476
1477 def test_init(self):
1478 class Planet(Enum):
1479 MERCURY = (3.303e+23, 2.4397e6)
1480 VENUS = (4.869e+24, 6.0518e6)
1481 EARTH = (5.976e+24, 6.37814e6)
1482 MARS = (6.421e+23, 3.3972e6)
1483 JUPITER = (1.9e+27, 7.1492e7)
1484 SATURN = (5.688e+26, 6.0268e7)
1485 URANUS = (8.686e+25, 2.5559e7)
1486 NEPTUNE = (1.024e+26, 2.4746e7)
1487 def __init__(self, mass, radius):
1488 self.mass = mass # in kilograms
1489 self.radius = radius # in meters
1490 @property
1491 def surface_gravity(self):
1492 # universal gravitational constant (m3 kg-1 s-2)
1493 G = 6.67300E-11
1494 return G * self.mass / (self.radius * self.radius)
1495 self.assertEqual(round(Planet.EARTH.surface_gravity, 2), 9.80)
1496 self.assertEqual(Planet.EARTH.value, (5.976e+24, 6.37814e6))
1497
Ethan Furman2aa27322013-07-19 19:35:56 -07001498 def test_nonhash_value(self):
1499 class AutoNumberInAList(Enum):
1500 def __new__(cls):
1501 value = [len(cls.__members__) + 1]
1502 obj = object.__new__(cls)
Ethan Furman520ad572013-07-19 19:47:21 -07001503 obj._value_ = value
Ethan Furman2aa27322013-07-19 19:35:56 -07001504 return obj
1505 class ColorInAList(AutoNumberInAList):
1506 red = ()
1507 green = ()
1508 blue = ()
1509 self.assertEqual(list(ColorInAList), [ColorInAList.red, ColorInAList.green, ColorInAList.blue])
Ethan Furman1a162882013-10-16 19:09:31 -07001510 for enum, value in zip(ColorInAList, range(3)):
1511 value += 1
1512 self.assertEqual(enum.value, [value])
1513 self.assertIs(ColorInAList([value]), enum)
Ethan Furman2aa27322013-07-19 19:35:56 -07001514
Ethan Furmanb41803e2013-07-25 13:50:45 -07001515 def test_conflicting_types_resolved_in_new(self):
1516 class LabelledIntEnum(int, Enum):
1517 def __new__(cls, *args):
1518 value, label = args
1519 obj = int.__new__(cls, value)
1520 obj.label = label
1521 obj._value_ = value
1522 return obj
1523
1524 class LabelledList(LabelledIntEnum):
1525 unprocessed = (1, "Unprocessed")
1526 payment_complete = (2, "Payment Complete")
1527
1528 self.assertEqual(list(LabelledList), [LabelledList.unprocessed, LabelledList.payment_complete])
1529 self.assertEqual(LabelledList.unprocessed, 1)
1530 self.assertEqual(LabelledList(1), LabelledList.unprocessed)
Ethan Furman2aa27322013-07-19 19:35:56 -07001531
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001532
Ethan Furmanf24bb352013-07-18 17:05:39 -07001533class TestUnique(unittest.TestCase):
1534
1535 def test_unique_clean(self):
1536 @unique
1537 class Clean(Enum):
1538 one = 1
1539 two = 'dos'
1540 tres = 4.0
1541 @unique
1542 class Cleaner(IntEnum):
1543 single = 1
1544 double = 2
1545 triple = 3
1546
1547 def test_unique_dirty(self):
1548 with self.assertRaisesRegex(ValueError, 'tres.*one'):
1549 @unique
1550 class Dirty(Enum):
1551 one = 1
1552 two = 'dos'
1553 tres = 1
1554 with self.assertRaisesRegex(
1555 ValueError,
1556 'double.*single.*turkey.*triple',
1557 ):
1558 @unique
1559 class Dirtier(IntEnum):
1560 single = 1
1561 double = 1
1562 triple = 3
1563 turkey = 3
1564
1565
Ethan Furman3323da92015-04-11 09:39:59 -07001566expected_help_output_with_docs = """\
Ethan Furman5875d742013-10-21 20:45:55 -07001567Help on class Color in module %s:
1568
1569class Color(enum.Enum)
Ethan Furman48a724f2015-04-11 23:23:06 -07001570 | An enumeration.
Serhiy Storchakab599ca82015-04-04 12:48:04 +03001571 |\x20\x20
Ethan Furman5875d742013-10-21 20:45:55 -07001572 | Method resolution order:
1573 | Color
1574 | enum.Enum
1575 | builtins.object
1576 |\x20\x20
1577 | Data and other attributes defined here:
1578 |\x20\x20
1579 | blue = <Color.blue: 3>
1580 |\x20\x20
1581 | green = <Color.green: 2>
1582 |\x20\x20
1583 | red = <Color.red: 1>
1584 |\x20\x20
1585 | ----------------------------------------------------------------------
1586 | Data descriptors inherited from enum.Enum:
1587 |\x20\x20
1588 | name
1589 | The name of the Enum member.
1590 |\x20\x20
1591 | value
1592 | The value of the Enum member.
1593 |\x20\x20
1594 | ----------------------------------------------------------------------
1595 | Data descriptors inherited from enum.EnumMeta:
1596 |\x20\x20
1597 | __members__
1598 | Returns a mapping of member name->value.
1599 |\x20\x20\x20\x20\x20\x20
1600 | This mapping lists all enum members, including aliases. Note that this
Ethan Furman3323da92015-04-11 09:39:59 -07001601 | is a read-only view of the internal mapping."""
1602
1603expected_help_output_without_docs = """\
1604Help on class Color in module %s:
1605
1606class Color(enum.Enum)
1607 | Method resolution order:
1608 | Color
1609 | enum.Enum
1610 | builtins.object
1611 |\x20\x20
1612 | Data and other attributes defined here:
1613 |\x20\x20
1614 | blue = <Color.blue: 3>
1615 |\x20\x20
1616 | green = <Color.green: 2>
1617 |\x20\x20
1618 | red = <Color.red: 1>
1619 |\x20\x20
1620 | ----------------------------------------------------------------------
1621 | Data descriptors inherited from enum.Enum:
1622 |\x20\x20
1623 | name
1624 |\x20\x20
1625 | value
1626 |\x20\x20
1627 | ----------------------------------------------------------------------
1628 | Data descriptors inherited from enum.EnumMeta:
1629 |\x20\x20
1630 | __members__"""
Ethan Furman5875d742013-10-21 20:45:55 -07001631
1632class TestStdLib(unittest.TestCase):
1633
Ethan Furman48a724f2015-04-11 23:23:06 -07001634 maxDiff = None
1635
Ethan Furman5875d742013-10-21 20:45:55 -07001636 class Color(Enum):
1637 red = 1
1638 green = 2
1639 blue = 3
1640
1641 def test_pydoc(self):
1642 # indirectly test __objclass__
Ethan Furman3323da92015-04-11 09:39:59 -07001643 if StrEnum.__doc__ is None:
1644 expected_text = expected_help_output_without_docs % __name__
1645 else:
1646 expected_text = expected_help_output_with_docs % __name__
Ethan Furman5875d742013-10-21 20:45:55 -07001647 output = StringIO()
1648 helper = pydoc.Helper(output=output)
1649 helper(self.Color)
1650 result = output.getvalue().strip()
Victor Stinner4b0432d2014-06-16 22:48:43 +02001651 self.assertEqual(result, expected_text)
Ethan Furman5875d742013-10-21 20:45:55 -07001652
1653 def test_inspect_getmembers(self):
1654 values = dict((
1655 ('__class__', EnumMeta),
Ethan Furman48a724f2015-04-11 23:23:06 -07001656 ('__doc__', 'An enumeration.'),
Ethan Furman5875d742013-10-21 20:45:55 -07001657 ('__members__', self.Color.__members__),
1658 ('__module__', __name__),
1659 ('blue', self.Color.blue),
1660 ('green', self.Color.green),
1661 ('name', Enum.__dict__['name']),
1662 ('red', self.Color.red),
1663 ('value', Enum.__dict__['value']),
1664 ))
1665 result = dict(inspect.getmembers(self.Color))
1666 self.assertEqual(values.keys(), result.keys())
1667 failed = False
1668 for k in values.keys():
1669 if result[k] != values[k]:
1670 print()
1671 print('\n%s\n key: %s\n result: %s\nexpected: %s\n%s\n' %
1672 ('=' * 75, k, result[k], values[k], '=' * 75), sep='')
1673 failed = True
1674 if failed:
1675 self.fail("result does not equal expected, see print above")
1676
1677 def test_inspect_classify_class_attrs(self):
1678 # indirectly test __objclass__
1679 from inspect import Attribute
1680 values = [
1681 Attribute(name='__class__', kind='data',
1682 defining_class=object, object=EnumMeta),
1683 Attribute(name='__doc__', kind='data',
Ethan Furman48a724f2015-04-11 23:23:06 -07001684 defining_class=self.Color, object='An enumeration.'),
Ethan Furman5875d742013-10-21 20:45:55 -07001685 Attribute(name='__members__', kind='property',
1686 defining_class=EnumMeta, object=EnumMeta.__members__),
1687 Attribute(name='__module__', kind='data',
1688 defining_class=self.Color, object=__name__),
1689 Attribute(name='blue', kind='data',
1690 defining_class=self.Color, object=self.Color.blue),
1691 Attribute(name='green', kind='data',
1692 defining_class=self.Color, object=self.Color.green),
1693 Attribute(name='red', kind='data',
1694 defining_class=self.Color, object=self.Color.red),
1695 Attribute(name='name', kind='data',
1696 defining_class=Enum, object=Enum.__dict__['name']),
1697 Attribute(name='value', kind='data',
1698 defining_class=Enum, object=Enum.__dict__['value']),
1699 ]
1700 values.sort(key=lambda item: item.name)
1701 result = list(inspect.classify_class_attrs(self.Color))
1702 result.sort(key=lambda item: item.name)
1703 failed = False
1704 for v, r in zip(values, result):
1705 if r != v:
1706 print('\n%s\n%s\n%s\n%s\n' % ('=' * 75, r, v, '=' * 75), sep='')
1707 failed = True
1708 if failed:
1709 self.fail("result does not equal expected, see print above")
1710
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001711if __name__ == '__main__':
1712 unittest.main()