From 6e0344a4423c9840776485372e16684eee5c9f0b Mon Sep 17 00:00:00 2001
From: Aarni Koskela <akx@iki.fi>
Date: Fri, 29 Apr 2016 15:45:17 +0300
Subject: [PATCH] Add IntEnum (like the standard library's IntEnum)

---
 enumfields/__init__.py    | 2 +-
 enumfields/enums.py       | 6 ++++++
 tests/models.py           | 5 +++--
 tests/test_enums.py       | 4 ++++
 tests/test_form_fields.py | 7 ++++---
 5 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/enumfields/__init__.py b/enumfields/__init__.py
index 5bbb9e9..35a023e 100644
--- a/enumfields/__init__.py
+++ b/enumfields/__init__.py
@@ -1,2 +1,2 @@
-from .enums import Enum
+from .enums import Enum, IntEnum
 from .fields import EnumField, EnumIntegerField
diff --git a/enumfields/enums.py b/enumfields/enums.py
index 8a63c8f..9e0b9ae 100644
--- a/enumfields/enums.py
+++ b/enumfields/enums.py
@@ -42,3 +42,9 @@ def __str__(self):
         Show our label when Django uses the Enum for displaying in a view
         """
         return force_text(self.label)
+
+
+@python_2_unicode_compatible
+class IntEnum(int, Enum):
+    def __str__(self):  # See Enum.__str__
+        return force_text(self.label)
diff --git a/tests/models.py b/tests/models.py
index 9598f92..2c82480 100644
--- a/tests/models.py
+++ b/tests/models.py
@@ -1,7 +1,6 @@
 from django.db import models
-from enum import Enum, IntEnum
 
-from enumfields import EnumField, EnumIntegerField
+from enumfields import Enum, IntEnum, EnumField, EnumIntegerField
 
 
 class MyModel(models.Model):
@@ -26,6 +25,8 @@ class ZeroEnum(Enum):
     class IntegerEnum(IntEnum):
         A = 0
         B = 1
+        class Labels:
+            A = 'foo'
 
     taste = EnumField(Taste, default=Taste.SWEET)
     taste_null_default = EnumField(Taste, null=True, blank=True, default=None)
diff --git a/tests/test_enums.py b/tests/test_enums.py
index eb6831e..c395399 100644
--- a/tests/test_enums.py
+++ b/tests/test_enums.py
@@ -8,6 +8,8 @@
 import six
 from six import u
 
+from tests.models import MyModel
+
 
 class Color(Enum):
     __order__ = 'RED GREEN BLUE'
@@ -35,11 +37,13 @@ def test_custom_labels():
     # Custom label
     assert Color.RED.label == 'Reddish'
     assert six.text_type(Color.RED) == 'Reddish'
+    assert six.text_type(MyModel.IntegerEnum.A) == 'foo'
 
 def test_automatic_labels():
     # Automatic label
     assert Color.GREEN.label == 'Green'
     assert six.text_type(Color.GREEN) == 'Green'
+    assert six.text_type(MyModel.IntegerEnum.B) == 'B'
 
 def test_lazy_labels():
     # Lazy label
diff --git a/tests/test_form_fields.py b/tests/test_form_fields.py
index 8e8cf72..9347e21 100644
--- a/tests/test_form_fields.py
+++ b/tests/test_form_fields.py
@@ -1,5 +1,5 @@
 # -- encoding: UTF-8 --
-
+from django.db.models import BLANK_CHOICE_DASH
 from django.forms.models import modelform_factory
 import pytest
 from .models import MyModel
@@ -8,7 +8,7 @@
 
 def get_form(**kwargs):
     instance = MyModel(color=MyModel.Color.RED)
-    FormClass = modelform_factory(MyModel, fields=("color", "zero2"))
+    FormClass = modelform_factory(MyModel, fields=("color", "zero2", "int_enum"))
     return FormClass(instance=instance, **kwargs)
 
 
@@ -26,4 +26,5 @@ def test_bound_form_with_instance():
 
 def test_choices():
     form = get_form()
-    assert form.base_fields["zero2"].choices == [(0, 'ZERO'), (1, 'ONE')]
+    assert form.base_fields["zero2"].choices == [(0, 'Zero'), (1, 'One')]
+    assert form.base_fields["int_enum"].choices == BLANK_CHOICE_DASH + [(0, 'foo'), (1, 'B')]