From a2a6b754befc9dbd61dc26db6d5deab2db180f52 Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <stretch@packetlife.net>
Date: Thu, 7 Nov 2019 09:32:25 -0500
Subject: [PATCH] Introduce ChoiceSet class for field choices

---
 netbox/utilities/api.py     | 15 ++++++++++++++-
 netbox/utilities/choices.py | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 1 deletion(-)
 create mode 100644 netbox/utilities/choices.py

diff --git a/netbox/utilities/api.py b/netbox/utilities/api.py
index 9adfd84ad..36e37667f 100644
--- a/netbox/utilities/api.py
+++ b/netbox/utilities/api.py
@@ -13,6 +13,7 @@ from rest_framework.response import Response
 from rest_framework.serializers import Field, ModelSerializer, ValidationError
 from rest_framework.viewsets import ModelViewSet as _ModelViewSet, ViewSet
 
+from utilities.choices import ChoiceSet
 from .utils import dict_to_filter_params, dynamic_import
 
 
@@ -64,14 +65,17 @@ class ChoiceField(Field):
     Represent a ChoiceField as {'value': <DB value>, 'label': <string>}.
     """
     def __init__(self, choices, **kwargs):
+        self.choiceset = choices
         self._choices = dict()
+
+        # Unpack grouped choices
         for k, v in choices:
-            # Unpack grouped choices
             if type(v) in [list, tuple]:
                 for k2, v2 in v:
                     self._choices[k2] = v2
             else:
                 self._choices[k] = v
+
         super().__init__(**kwargs)
 
     def to_representation(self, obj):
@@ -81,6 +85,11 @@ class ChoiceField(Field):
             ('value', obj),
             ('label', self._choices[obj])
         ])
+
+        # Include legacy numeric ID (where applicable)
+        if type(self.choiceset) is ChoiceSet and obj in self.choiceset.LEGACY_MAP:
+            data['id'] = self.choiceset.LEGACY_MAP.get(obj)
+
         return data
 
     def to_internal_value(self, data):
@@ -104,6 +113,10 @@ class ChoiceField(Field):
         try:
             if data in self._choices:
                 return data
+            # Check if data is a legacy numeric ID
+            slug = self.choiceset.id_to_slug(data)
+            if slug is not None:
+                return slug
         except TypeError:  # Input is an unhashable type
             pass
 
diff --git a/netbox/utilities/choices.py b/netbox/utilities/choices.py
new file mode 100644
index 000000000..367451dd5
--- /dev/null
+++ b/netbox/utilities/choices.py
@@ -0,0 +1,36 @@
+class ChoiceSetMeta(type):
+    """
+    Metaclass for ChoiceSet
+    """
+    def __call__(cls, *args, **kwargs):
+        # Django will check if a choices value is callable, and if so assume that it returns an iterable
+        return getattr(cls, 'CHOICES', ())
+
+    def __iter__(cls):
+        choices = getattr(cls, 'CHOICES', ())
+        return iter(choices)
+
+
+class ChoiceSet(metaclass=ChoiceSetMeta):
+
+    CHOICES = list()
+    LEGACY_MAP = dict()
+
+    @classmethod
+    def slug_to_id(cls, slug):
+        """
+        Return the legacy integer value corresponding to a slug.
+        """
+        return cls.LEGACY_MAP.get(slug)
+
+    @classmethod
+    def id_to_slug(cls, legacy_id):
+        """
+        Return the slug value corresponding to a legacy integer value.
+        """
+        if legacy_id in cls.LEGACY_MAP.values():
+            # Invert the legacy map to allow lookup by integer
+            legacy_map = dict([
+                (id, slug) for slug, id in cls.LEGACY_MAP.items()
+            ])
+            return legacy_map.get(legacy_id)