Repository: ipython/traitlets
Pull request #200 (2016-03-26 10:11:40)

Provide UseEnum as Traitlets class to support enum34 Enum classes 
CLEAN-REPLACEMENT FOR: pull #169

Add support for Enum classes (py34: enum.Enum) as model for the data type description (instead of string-based enums).

EXAMPLE:
```python
# -- SINCE: Python 3.4 (or install backport: pip install enum34)
import enum
from traitlets import HasTraits, UseEnum

class Color(enum.Enum):
    red = 1         # -- IMPLICIT: default_value
    blue = 2
    green = 3

class MyEntity(HasTraits):
    color = UseEnum(Color, default_value=Color.blue)

entity = MyEntity(color=Color.red)
entity.color = Color.green    # USE: Enum-value (preferred)
entity.color = "green"        # USE: name (as string)
entity.color = "Color.green"  # USE: scoped-name (as string)
entity.color = 3              # USE: number (as int)
assert entity.color is Color.green
```



Patch:
diff --git a/docs/source/trait_types.rst b/docs/source/trait_types.rst
index 96beab5..3c67003 100644
--- a/docs/source/trait_types.rst
+++ b/docs/source/trait_types.rst
@@ -101,6 +101,8 @@ Miscellaneous
 
 .. autoclass:: CaselessStrEnum
 
+.. autoclass:: UseEnum
+
 .. autoclass:: TCPAddress
 
 .. autoclass:: CRegExp
diff --git a/setup.py b/setup.py
index 0a15aec..2d77eb5 100644
--- a/setup.py
+++ b/setup.py
@@ -83,6 +83,10 @@ install_requires = setuptools_args['install_requires'] = [
 
 extras_require = setuptools_args['extras_require'] = {
     'test': ['nose2'],
+    # -- SUPPORT UNIFORM-WHEELS: Extra packages for Python 2.7, 3.3
+    # SEE: https://bitbucket.org/pypa/wheel/ , CHANGES.txt (v0.24.0)
+    ':python_version=="2.7"': ["enum34"],
+    ':python_version=="3.3"': ["enum34"],
 }
 
 if 'setuptools' in sys.modules:
diff --git a/traitlets/traitlets.py b/traitlets/traitlets.py
index 080264c..de4f673 100644
--- a/traitlets/traitlets.py
+++ b/traitlets/traitlets.py
@@ -45,6 +45,7 @@ import inspect
 import re
 import sys
 import types
+import enum
 try:
     from types import ClassType, InstanceType
     ClassTypes = (ClassType, type)
@@ -2471,3 +2472,88 @@ class CRegExp(TraitType):
             return re.compile(value)
         except:
             self.error(obj, value)
+
+
+class UseEnum(TraitType):
+    """Use a Enum class as model for the data type description.
+    Note that if no default-value is provided, the first enum-value is used
+    as default-value.
+
+    .. sourcecode:: python
+
+        # -- SINCE: Python 3.4 (or install backport: pip install enum34)
+        import enum
+        from traitlets import HasTraits, UseEnum
+
+        class Color(enum.Enum):
+            red = 1         # -- IMPLICIT: default_value
+            blue = 2
+            green = 3
+
+        class MyEntity(HasTraits):
+            color = UseEnum(Color, default_value=Color.blue)
+
+        entity = MyEntity(color=Color.red)
+        entity.color = Color.green    # USE: Enum-value (preferred)
+        entity.color = "green"        # USE: name (as string)
+        entity.color = "Color.green"  # USE: scoped-name (as string)
+        entity.color = 3              # USE: number (as int)
+        assert entity.color is Color.green
+    """
+    default_value = None
+    info_text = "Trait type adapter to a Enum class"
+
+    def __init__(self, enum_class, default_value=None, **kwargs):
+        assert issubclass(enum_class, enum.Enum), \
+                          "REQUIRE: enum.Enum, but was: %r" % enum_class
+        allow_none = kwargs.get("allow_none", False)
+        if default_value is None and not allow_none:
+            default_value = list(enum_class.__members__.values())[0]
+        super(UseEnum, self).__init__(default_value=default_value, **kwargs)
+        self.enum_class = enum_class
+        self.name_prefix = enum_class.__name__ + "."
+
+    def select_by_number(self, value, default=Undefined):
+        """Selects enum-value by using its number-constant."""
+        assert isinstance(value, int)
+        enum_members = self.enum_class.__members__
+        for enum_item in enum_members.values():
+            if enum_item.value == value:
+                return enum_item
+        # -- NOT FOUND:
+        return default
+
+    def select_by_name(self, value, default=Undefined):
+        """Selects enum-value by using its name or scoped-name."""
+        assert isinstance(value, string_types)
+        if value.startswith(self.name_prefix):
+            # -- SUPPORT SCOPED-NAMES, like: "Color.red" => "red"
+            value = value.replace(self.name_prefix, "", 1)
+        return self.enum_class.__members__.get(value, default)
+
+    def validate(self, obj, value):
+        if isinstance(value, self.enum_class):
+            return value
+        elif isinstance(value, int):
+            # -- CONVERT: number => enum_value (item)
+            value2 = self.select_by_number(value)
+            if value2 is not Undefined:
+                return value2
+        elif isinstance(value, string_types):
+            # -- CONVERT: name or scoped_name (as string) => enum_value (item)
+            value2 = self.select_by_name(value)
+            if value2 is not Undefined:
+                return value2
+        elif value is None:
+            if self.allow_none:
+                return None
+            else:
+                return self.default_value
+        self.error(obj, value)
+
+    def info(self):
+        """Returns a description of this Enum trait (in case of errors)."""
+        result = "Any of: %s" % ", ".join(self.enum_class.__members__.keys())
+        if self.allow_none:
+            return result + " or None"
+        return result


Test patch:
diff --git a/traitlets/tests/test_traitlets_enum.py b/traitlets/tests/test_traitlets_enum.py
new file mode 100644
index 0000000..82259ae
--- /dev/null
+++ b/traitlets/tests/test_traitlets_enum.py
@@ -0,0 +1,181 @@
+# -*- coding: UTF-8 -*-
+# pylint: disable=missing-docstring, too-few-public-methods
+"""
+Test the trait-type ``UseEnum``.
+"""
+
+import unittest
+import enum
+from ipython_genutils.py3compat import string_types
+from traitlets import HasTraits, TraitError, UseEnum
+
+
+# -----------------------------------------------------------------------------
+# TEST SUPPORT:
+# -----------------------------------------------------------------------------
+class Color(enum.Enum):
+    red = 1
+    green = 2
+    blue = 3
+    yellow = 4
+
+class OtherColor(enum.Enum):
+    red = 0
+    green = 1
+
+
+# -----------------------------------------------------------------------------
+# TESTSUITE:
+# -----------------------------------------------------------------------------
+class TestUseEnum(unittest.TestCase):
+    # pylint: disable=invalid-name
+
+    class Example(HasTraits):
+        color = UseEnum(Color, help="Color enum")
+
+    def test_assign_enum_value(self):
+        example = self.Example()
+        example.color = Color.green
+        self.assertEqual(example.color, Color.green)
+
+    def test_assign_all_enum_values(self):
+        # pylint: disable=no-member
+        enum_values = [value  for value in Color.__members__.values()]
+        for value in enum_values:
+            self.assertIsInstance(value, Color)
+            example = self.Example()
+            example.color = value
+            self.assertEqual(example.color, value)
+            self.assertIsInstance(value, Color)
+
+    def test_assign_enum_value__with_other_enum_raises_error(self):
+        example = self.Example()
+        with self.assertRaises(TraitError):
+            example.color = OtherColor.green
+
+    def test_assign_enum_name_1(self):
+        # -- CONVERT: string => Enum value (item)
+        example = self.Example()
+        example.color = "red"
+        self.assertEqual(example.color, Color.red)
+
+    def test_assign_enum_value_name(self):
+        # -- CONVERT: string => Enum value (item)
+        # pylint: disable=no-member
+        enum_names = [enum_val.name for enum_val in Color.__members__.values()]
+        for value in enum_names:
+            self.assertIsInstance(value, string_types)
+            example = self.Example()
+            enum_value = Color.__members__.get(value)
+            example.color = value
+            self.assertIs(example.color, enum_value)
+            self.assertEqual(example.color.name, value)
+
+    def test_assign_scoped_enum_value_name(self):
+        # -- CONVERT: string => Enum value (item)
+        scoped_names = ["Color.red", "Color.green", "Color.blue", "Color.yellow"]
+        for value in scoped_names:
+            example = self.Example()
+            example.color = value
+            self.assertIsInstance(example.color, Color)
+            self.assertEqual(str(example.color), value)
+
+    def test_assign_bad_enum_value_name__raises_error(self):
+        # -- CONVERT: string => Enum value (item)
+        bad_enum_names = ["UNKNOWN_COLOR", "RED", "Green", "blue2"]
+        for value in bad_enum_names:
+            example = self.Example()
+            with self.assertRaises(TraitError):
+                example.color = value
+
+    def test_assign_enum_value_number_1(self):
+        # -- CONVERT: number => Enum value (item)
+        example = self.Example()
+        example.color = 1  # == Color.red.value
+        example.color = Color.red.value
+        self.assertEqual(example.color, Color.red)
+
+    def test_assign_enum_value_number(self):
+        # -- CONVERT: number => Enum value (item)
+        # pylint: disable=no-member
+        enum_numbers = [enum_val.value
+                        for enum_val in Color.__members__.values()]
+        for value in enum_numbers:
+            self.assertIsInstance(value, int)
+            example = self.Example()
+            example.color = value
+            self.assertIsInstance(example.color, Color)
+            self.assertEqual(example.color.value, value)
+
+    def test_assign_bad_enum_value_number__raises_error(self):
+        # -- CONVERT: number => Enum value (item)
+        bad_numbers = [-1, 0, 5]
+        for value in bad_numbers:
+            self.assertIsInstance(value, int)
+            assert UseEnum(Color).select_by_number(value, None) is None
+            example = self.Example()
+            with self.assertRaises(TraitError):
+                example.color = value
+
+    def test_ctor_without_default_value(self):
+        # -- IMPLICIT: default_value = Color.red (first enum-value)
+        class Example2(HasTraits):
+            color = UseEnum(Color)
+
+        example = Example2()
+        self.assertEqual(example.color, Color.red)
+
+    def test_ctor_with_default_value_as_enum_value(self):
+        # -- CONVERT: number => Enum value (item)
+        class Example2(HasTraits):
+            color = UseEnum(Color, default_value=Color.green)
+
+        example = Example2()
+        self.assertEqual(example.color, Color.green)
+
+
+    def test_ctor_with_default_value_none_and_not_allow_none(self):
+        # -- IMPLICIT: default_value = Color.red (first enum-value)
+        class Example2(HasTraits):
+            color1 = UseEnum(Color, default_value=None, allow_none=False)
+            color2 = UseEnum(Color, default_value=None)
+        example = Example2()
+        self.assertEqual(example.color1, Color.red)
+        self.assertEqual(example.color2, Color.red)
+
+    def test_ctor_with_default_value_none_and_allow_none(self):
+        class Example2(HasTraits):
+            color1 = UseEnum(Color, default_value=None, allow_none=True)
+            color2 = UseEnum(Color, allow_none=True)
+
+        example = Example2()
+        self.assertIs(example.color1, None)
+        self.assertIs(example.color2, None)
+
+    def test_assign_none_without_allow_none_resets_to_default_value(self):
+        class Example2(HasTraits):
+            color1 = UseEnum(Color, allow_none=False)
+            color2 = UseEnum(Color)
+
+        example = Example2()
+        example.color1 = None
+        example.color2 = None
+        self.assertIs(example.color1, Color.red)
+        self.assertIs(example.color2, Color.red)
+
+    def test_assign_none_to_enum_or_none(self):
+        class Example2(HasTraits):
+            color = UseEnum(Color, allow_none=True)
+
+        example = Example2()
+        example.color = None
+        self.assertIs(example.color, None)
+
+    def test_assign_bad_value_with_to_enum_or_none(self):
+        class Example2(HasTraits):
+            color = UseEnum(Color, allow_none=True)
+
+        example = Example2()
+        with self.assertRaises(TraitError):
+            example.color = "BAD_VALUE"
+