1 files changed, 24 insertions, 1 deletions
diff --git a/ihatemoney/utils.py b/ihatemoney/utils.py
index 126b9de..0641d1c 100644
--- a/ihatemoney/utils.py
+++ b/ihatemoney/utils.py
@@ -2,6 +2,7 @@ import re
 import os
 import ast
 import operator
+from enum import Enum
 
 from io import BytesIO, StringIO
 
@@ -12,7 +13,6 @@ from babel import Locale
 from werkzeug.routing import HTTPException, RoutingException
 from datetime import datetime, timedelta
 
-
 import csv
 
 
@@ -257,3 +257,26 @@ def same_bill(bill1, bill2):
         if bill1[a] != bill2[a]:
             return False
     return True
+
+
+class FormEnum(Enum):
+    """Extend builtin Enum class to be seamlessly compatible with WTForms"""
+
+    @classmethod
+    def choices(cls):
+        return [(choice, choice.name) for choice in cls]
+
+    @classmethod
+    def coerce(cls, item):
+        """Coerce a str or int representation into an Enum object"""
+        if isinstance(item, cls):
+            return item
+
+        # If item is not already a Enum object then it must be
+        # a string or int corresponding to an ID (e.g. '0' or 1)
+        # Either int() or cls() will correctly throw a TypeError if this
+        # is not the case
+        return cls(int(item))
+
+    def __str__(self):
+        return str(self.value)