From 24f30d46b25f9aebe3ed833454fd8d074dbbecd1 Mon Sep 17 00:00:00 2001
From: Peter Hamilton <peter.allen.hamilton@gmail.com>
Date: Mon, 25 Feb 2019 10:48:50 -0500
Subject: [PATCH] Add OrderedEnum to enable enum value comparisons

This change adds a new Enum subclass, OrderedEnum, that provides
comparison operators allowing for comparisons between enum values.
This will be used going forward with the KMIPVersion enum enabling
version checking on supported or unsupported object types.
---
 kmip/core/enums.py                 | 30 ++++++++++++++-
 kmip/tests/unit/core/test_enums.py | 61 ++++++++++++++++++++++++++++++
 2 files changed, 90 insertions(+), 1 deletion(-)

diff --git a/kmip/core/enums.py b/kmip/core/enums.py
index 0f155c1..fb3fdce 100644
--- a/kmip/core/enums.py
+++ b/kmip/core/enums.py
@@ -20,6 +20,34 @@ import copy
 import enum
 
 
+class OrderedEnum(enum.Enum):
+    """
+    An ordered variant of the Enum class that allows for comparisons.
+
+    Taken from: https://docs.python.org/3/library/enum.html#orderedenum
+    """
+
+    def __ge__(self, other):
+        if self.__class__ is other.__class__:
+            return self.value >= other.value
+        return NotImplemented
+
+    def __gt__(self, other):
+        if self.__class__ is other.__class__:
+            return self.value > other.value
+        return NotImplemented
+
+    def __le__(self, other):
+        if self.__class__ is other.__class__:
+            return self.value <= other.value
+        return NotImplemented
+
+    def __lt__(self, other):
+        if self.__class__ is other.__class__:
+            return self.value < other.value
+        return NotImplemented
+
+
 class AdjustmentType(enum.Enum):
     # KMIP 2.0
     INCREMENT = 0x00000001
@@ -508,7 +536,7 @@ class KeyWrapType(enum.Enum):
     AS_REGISTERED = 0x00000002
 
 
-class KMIPVersion(enum.Enum):
+class KMIPVersion(OrderedEnum):
     KMIP_1_0 = 1.0
     KMIP_1_1 = 1.1
     KMIP_1_2 = 1.2
diff --git a/kmip/tests/unit/core/test_enums.py b/kmip/tests/unit/core/test_enums.py
index 6d9b3e6..517af20 100644
--- a/kmip/tests/unit/core/test_enums.py
+++ b/kmip/tests/unit/core/test_enums.py
@@ -18,6 +18,67 @@ import testtools
 from kmip.core import enums
 
 
+class TestOrderedEnum(testtools.TestCase):
+
+    def setUp(self):
+        super(TestOrderedEnum, self).setUp()
+
+    def tearDown(self):
+        super(TestOrderedEnum, self).tearDown()
+
+    def test_greater_than_or_equal(self):
+        self.assertTrue(
+            enums.KMIPVersion.KMIP_2_0 >= enums.KMIPVersion.KMIP_1_0
+        )
+        self.assertFalse(
+            enums.KMIPVersion.KMIP_1_0 >= enums.KMIPVersion.KMIP_2_0
+        )
+
+        self.assertEquals(
+            NotImplemented,
+            enums.KMIPVersion.KMIP_2_0.__ge__(enums.WrappingMethod.ENCRYPT)
+        )
+
+    def test_greater_than(self):
+        self.assertTrue(
+            enums.KMIPVersion.KMIP_1_3 > enums.KMIPVersion.KMIP_1_1
+        )
+        self.assertFalse(
+            enums.KMIPVersion.KMIP_1_1 > enums.KMIPVersion.KMIP_1_3
+        )
+
+        self.assertEquals(
+            NotImplemented,
+            enums.KMIPVersion.KMIP_2_0.__gt__(enums.WrappingMethod.ENCRYPT)
+        )
+
+    def test_less_than_or_equal(self):
+        self.assertTrue(
+            enums.KMIPVersion.KMIP_1_3 <= enums.KMIPVersion.KMIP_1_4
+        )
+        self.assertFalse(
+            enums.KMIPVersion.KMIP_1_4 <= enums.KMIPVersion.KMIP_1_3
+        )
+
+        self.assertEquals(
+            NotImplemented,
+            enums.KMIPVersion.KMIP_2_0.__le__(enums.WrappingMethod.ENCRYPT)
+        )
+
+    def test_less_than(self):
+        self.assertTrue(
+            enums.KMIPVersion.KMIP_1_3 < enums.KMIPVersion.KMIP_2_0
+        )
+        self.assertFalse(
+            enums.KMIPVersion.KMIP_2_0 < enums.KMIPVersion.KMIP_1_3
+        )
+
+        self.assertEquals(
+            NotImplemented,
+            enums.KMIPVersion.KMIP_2_0.__lt__(enums.WrappingMethod.ENCRYPT)
+        )
+
+
 class TestEnumUtilityFunctions(testtools.TestCase):
 
     def setUp(self):