diff --git a/voluptuous/tests/tests.py b/voluptuous/tests/tests.py index cfac3e7..d87f301 100644 --- a/voluptuous/tests/tests.py +++ b/voluptuous/tests/tests.py @@ -1,5 +1,6 @@ import copy import collections +from enum import Enum import os import sys @@ -1599,3 +1600,39 @@ def test_any_with_discriminant(): assert_equal(str(e), 'expected bool for dictionary value @ data[\'implementation\'][\'c-value\']') else: assert False, "Did not raise correct Invalid" + +def test_coerce_enum(): + """Test Coerce Enum""" + class Choice(Enum): + Easy = 1 + Medium = 2 + Hard = 3 + + class StringChoice(str, Enum): + Easy = "easy" + Medium = "medium" + Hard = "hard" + + schema = Schema(Coerce(Choice)) + string_schema = Schema(Coerce(StringChoice)) + + # Valid value + assert schema(1) == Choice.Easy + assert string_schema("easy") == StringChoice.Easy + + # Invalid value + try: + schema(4) + except Invalid as e: + assert_equal(str(e), + "expected Choice or one of 1, 2, 3") + else: + assert False, "Did not raise Invalid for String" + + try: + string_schema("hello") + except Invalid as e: + assert_equal(str(e), + "expected StringChoice or one of 'easy', 'medium', 'hard'") + else: + assert False, "Did not raise Invalid for String" diff --git a/voluptuous/validators.py b/voluptuous/validators.py index 5f72cf6..972fdde 100644 --- a/voluptuous/validators.py +++ b/voluptuous/validators.py @@ -4,6 +4,7 @@ import sys from functools import wraps from decimal import Decimal, InvalidOperation +from enum import Enum from voluptuous.schema_builder import Schema, raises, message from voluptuous.error import (MultipleInvalid, CoerceInvalid, TrueInvalid, FalseInvalid, BooleanInvalid, Invalid, @@ -95,6 +96,8 @@ def __call__(self, v): return self.type(v) except (ValueError, TypeError, InvalidOperation): msg = self.msg or ('expected %s' % self.type_name) + if not self.msg and issubclass(self.type, Enum): + msg += " or one of %s" % str([e.value for e in self.type])[1:-1] raise CoerceInvalid(msg) def __repr__(self):