Introduce EnumExtension metaclass

This metaclass allows extensions to register new values in any
enum by (ab)using metaclasses.

For example, this allows for extensions to introduce new
RequirementFlags without having to worry about import order.

gssapi/raw/misc.pyx

@@ -11,6 +11,8 @@ from gssapi.raw.oids cimport OID
 
 from gssapi.raw.types import MechType
 
+from enum import EnumMeta, Enum
+
 
 cdef extern from "python_gssapi.h":
     OM_uint32 gss_display_status(OM_uint32 *minor_status,
@@ -329,3 +331,64 @@ class GSSError(Exception):
                                    maj_str=maj_str,
                                    min_stat=self.min_code,
                                    min_str=min_str)
+
+
+# WARNING: FOR YOUR OWN PERSONAL SANITY, DO NOT TRY THIS AT HOME
+# This allows things like extensions declaring their own RequirementFlags
+# without having to worry about import order.  There *might* be a cleaner
+# way to do this, but most of the ways probably exploit the interals of Enum
+class _EnumExtension(EnumMeta):
+    def __new__(metacls, name, bases, classdict):
+        # find the base class that this overrides
+
+        base_attr = classdict.get('__base__', None)
+
+        if len(bases) != 1 or bases[0] is not object or base_attr is None:
+            raise TypeError("Enumeration extensions must be created as "
+                            "`ClassName(object)` and must contain the "
+                            "`__base__` attribute")
+
+        adds_to = base_attr
+        del classdict['__base__']
+
+        bases = adds_to.__bases__
+
+        # we need to have the same Enum type
+        if not issubclass(adds_to, Enum):
+            raise TypeError("Enumeration extensions must extend a subtype "
+                            "of Enum or directly sublcass object")
+
+        # roughly preserve insertion order in Python 3.4+
+        # (actually preserving order would require using parts of the Enum
+        # implementation that aren't part of the spec)
+        old_classdict = classdict
+        classdict = metacls.__prepare__(name, bases)
+
+        # NB(directxman12): we can't use update, because that doesn't
+        # trigger __setitem__, which thus won't update the _EnumDict correctly
+        base_members = adds_to.__members__
+        for k, v in base_members.items():
+            classdict[k] = v.value
+
+        for k, v in old_classdict.items():
+            if k in base_members:
+                raise AttributeError("Enumeration extensions cannot override "
+                                     "existing enumeration members")
+
+            classdict[k] = v
+
+        res = super(_EnumExtension, metacls).__new__(metacls, name,
+                                                     bases, classdict)
+
+        # replace the old with the new
+        for k, v in res.__dict__.items():
+            if (k not in ('__module__', '__qualname__', '__doc__') and
+                    k in adds_to.__dict__):
+                setattr(adds_to, k, v)
+
+        # preserve enum semantics
+        for member in adds_to.__members__.values():
+            member.__class__ = adds_to
+
+        # always return the original
+        return adds_to

gssapi/tests/test_raw.py

@@ -5,6 +5,8 @@
 import unittest
 
 import should_be.all  # noqa
+from enum import IntEnum
+import six
 
 import gssapi.raw as gb
 import gssapi.raw.misc as gbmisc
@@ -641,6 +643,78 @@ def test_xor_set(self):
         fset3.should_include(gb.RequirementFlag.out_of_sequence_detection)
 
 
+class TestEnumExtension(unittest.TestCase):
+    def setUp(self):
+        class ExtendedEnum(IntEnum):
+            """some docs"""
+            a = 1
+            b = 2
+            c = 3
+
+        class OtherEnum(IntEnum):
+            a = 1
+            b = 2
+            c = 3
+
+        self._ext_enum = ExtendedEnum
+        self._plain_enum = OtherEnum
+
+    def test_same_dict(self):
+        @six.add_metaclass(gbmisc._EnumExtension)
+        class ExtendedEnum(object):
+            __base__ = self._ext_enum
+            d = 4
+            e = 5
+
+        self._ext_enum.__dict__.keys().should_be(
+                self._plain_enum.__dict__.keys())
+
+    def test_members_copy_mult(self):
+        @six.add_metaclass(gbmisc._EnumExtension)
+        class ExtendedEnum(object):
+            __base__ = self._ext_enum
+            d = 4
+
+        @six.add_metaclass(gbmisc._EnumExtension)
+        class ExtendedEnum2(object):
+            __base__ = self._ext_enum
+            e = 5
+
+        self._ext_enum.__members__.should_include(
+            ['a', 'b', 'c', 'd', 'e'])
+
+        self._ext_enum.a.should_be(1)
+        self._ext_enum.b.should_be(2)
+        self._ext_enum.c.should_be(3)
+        self._ext_enum.d.should_be(4)
+        self._ext_enum.e.should_be(5)
+
+    def test_all_class_are_same(self):
+        @six.add_metaclass(gbmisc._EnumExtension)
+        class ExtendedEnum(object):
+            __base__ = self._ext_enum
+            d = 4
+
+        assert ExtendedEnum is self._ext_enum
+
+    def test_members_are_still_instance(self):
+        @six.add_metaclass(gbmisc._EnumExtension)
+        class EnumExt(object):
+            __base__ = self._ext_enum
+            d = 4
+
+        self._ext_enum.a.should_be_a(self._ext_enum)
+        self._ext_enum.d.should_be_a(self._ext_enum)
+
+    def test_doc_is_preserved(self):
+        @six.add_metaclass(gbmisc._EnumExtension)
+        class ExtendedEnum(object):
+            __base__ = self._ext_enum
+            d = 4
+
+        self._ext_enum.__doc__.should_be('some docs')
+
+
 class TestInitContext(_GSSAPIKerberosTestCase):
     def setUp(self):
         self.target_name = gb.import_name(TARGET_SERVICE_NAME,