From 53e09a924c0d8331391ed02ea41aaca91c46ed2f Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <stretch@packetlife.net>
Date: Wed, 26 Aug 2020 15:04:22 -0400
Subject: [PATCH] Restore and rename CustomChoiceFieldInspector

---
 netbox/netbox/settings.py             |  1 +
 netbox/utilities/custom_inspectors.py | 37 +++++++++++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/netbox/netbox/settings.py b/netbox/netbox/settings.py
index 7bcc806d7..66aabfcf5 100644
--- a/netbox/netbox/settings.py
+++ b/netbox/netbox/settings.py
@@ -487,6 +487,7 @@ SWAGGER_SETTINGS = {
         'utilities.custom_inspectors.CustomFieldsDataFieldInspector',
         'utilities.custom_inspectors.JSONFieldInspector',
         'utilities.custom_inspectors.NullableBooleanFieldInspector',
+        'utilities.custom_inspectors.ChoiceFieldInspector',
         'utilities.custom_inspectors.SerializedPKRelatedFieldInspector',
         'drf_yasg.inspectors.CamelCaseJSONFilter',
         'drf_yasg.inspectors.ReferencingSerializerInspector',
diff --git a/netbox/utilities/custom_inspectors.py b/netbox/utilities/custom_inspectors.py
index bee2f7d92..1d5c9c0a0 100644
--- a/netbox/utilities/custom_inspectors.py
+++ b/netbox/utilities/custom_inspectors.py
@@ -49,6 +49,43 @@ class SerializedPKRelatedFieldInspector(FieldInspector):
         return NotHandled
 
 
+class ChoiceFieldInspector(FieldInspector):
+    def field_to_swagger_object(self, field, swagger_object_type, use_references, **kwargs):
+        # this returns a callable which extracts title, description and other stuff
+        # https://drf-yasg.readthedocs.io/en/stable/_modules/drf_yasg/inspectors/base.html#FieldInspector._get_partial_types
+        SwaggerType, _ = self._get_partial_types(field, swagger_object_type, use_references, **kwargs)
+
+        if isinstance(field, ChoiceField):
+            choices = field._choices
+            choice_value = list(choices.keys())
+            choice_label = list(choices.values())
+            value_schema = openapi.Schema(type=openapi.TYPE_STRING, enum=choice_value)
+
+            if set([None] + choice_value) == {None, True, False}:
+                # DeviceType.subdevice_role, Device.face and InterfaceConnection.connection_status all need to be
+                # differentiated since they each have subtly different values in their choice keys.
+                # - subdevice_role and connection_status are booleans, although subdevice_role includes None
+                # - face is an integer set {0, 1} which is easily confused with {False, True}
+                schema_type = openapi.TYPE_STRING
+                if all(type(x) == bool for x in [c for c in choice_value if c is not None]):
+                    schema_type = openapi.TYPE_BOOLEAN
+                value_schema = openapi.Schema(type=schema_type, enum=choice_value)
+                value_schema['x-nullable'] = True
+
+            if all(type(x) == int for x in [c for c in choice_value if c is not None]):
+                # Change value_schema for IPAddressFamilyChoices, RackWidthChoices
+                value_schema = openapi.Schema(type=openapi.TYPE_INTEGER, enum=choice_value)
+
+            schema = SwaggerType(type=openapi.TYPE_OBJECT, required=["label", "value"], properties={
+                "label": openapi.Schema(type=openapi.TYPE_STRING, enum=choice_label),
+                "value": value_schema
+            })
+
+            return schema
+
+        return NotHandled
+
+
 class NullableBooleanFieldInspector(FieldInspector):
     def process_result(self, result, method_name, obj, **kwargs):