trafaret/Enum: support Python >=3.4 enum module

Author: spumer

tests/test_base.py

@@ -1,6 +1,8 @@
 # -*- coding: utf-8 -*-
+import sys
 import unittest
 import trafaret as t
+import trafaret.utils as tu
 from collections import Mapping as AbcMapping
 from trafaret import extract_error, ignore, DataError
 from trafaret.extras import KeysSubset
@@ -264,6 +266,40 @@ def test_enum(self):
         res = extract_error(trafaret, 2)
         self.assertEqual(res, "value doesn't match any variant")
 
+    @unittest.skipIf(sys.version_info < (3, 4),
+        "not supported in this veresion"
+    )
+    def test_enum_py3(self):
+        import enum
+
+        class Colors(enum.Enum):
+            red = 0
+            green = 1
+            blue = 2
+
+        trafaret = t.Enum(Colors)
+        self.assertEqual(repr(trafaret), "<Enum('red', 'green', 'blue')>")
+        res = trafaret.check('red')
+        res = trafaret.check('green')
+        res = extract_error(trafaret, 'unknown')
+        self.assertEqual(res, "value doesn't match any variant")
+
+        # check multiple fails
+
+        class Fruits(enum.Enum):
+            orange = 0
+            apple = 1
+
+        with self.assertRaises(TypeError):
+            trafaret = t.Enum(Colors, Fruits)
+
+        # check mixin fails in any order
+
+        with self.assertRaises(TypeError):
+            trafaret = t.Enum(Colors, 1)
+
+        with self.assertRaises(TypeError):
+            trafaret = t.Enum(1, Colors)
 
 
 class TestFloat(unittest.TestCase):

trafaret/__init__.py

@@ -15,11 +15,19 @@
 __VERSION__ = (0, 10, 2)
 
 
+enum = None
+
 # Python3 support
 if py3:
     import urllib.parse as urlparse
     str_types = (str, bytes)
     unicode = str
+
+    try:
+        import enum
+    except ImportError:
+        pass
+
 else:
     try:
         from future_builtins import map
@@ -1345,6 +1353,14 @@ class Enum(Trafaret):
     __slots__ = ['variants']
 
     def __init__(self, *variants):
+        if enum is not None:
+            has_meta = any(isinstance(v, enum.EnumMeta) for v in variants)
+            if has_meta:
+                if len(variants) != 1:
+                    raise TypeError("You can't use enum.Enum with other arguments")
+
+                variants = [x.name for x in variants[0]]
+
         self.variants = variants[:]
 
     def check_value(self, value):