Find better way for fix enum problem

src/drf_yasg/codecs.py

@@ -113,12 +113,6 @@ def __init__(self, validators, pretty=False, media_type='application/json'):
         self.pretty = pretty
         self.media_type = media_type
 
-    @staticmethod
-    def defaults(value):
-        if isinstance(value, enum.Enum):
-            return value.value
-        raise TypeError
-
     def _dump_dict(self, spec):
         """Dump ``spec`` into JSON.
 
@@ -186,16 +180,11 @@ def represent_text(self, text):
             return self.represent_scalar('tag:yaml.org,2002:str', text, style='|')
         return self.represent_scalar('tag:yaml.org,2002:str', text)
 
-    def represent_enum(self, element: enum.Enum):
-        return self.represent_text(element.value)
-
 
 SaneYamlDumper.add_representer(bytes, SaneYamlDumper.represent_text)
 SaneYamlDumper.add_representer(str, SaneYamlDumper.represent_text)
 SaneYamlDumper.add_representer(OrderedDict, SaneYamlDumper.represent_odict)
-SaneYamlDumper.add_representer(enum.Enum, SaneYamlDumper.represent_enum)
 SaneYamlDumper.add_multi_representer(OrderedDict, SaneYamlDumper.represent_odict)
-SaneYamlDumper.add_multi_representer(enum.Enum, SaneYamlDumper.represent_enum)
 
 
 def yaml_sane_dump(data, binary):

src/drf_yasg/openapi.py

@@ -1,6 +1,7 @@
 import collections
 import logging
 import re
+import enum
 import urllib.parse as urlparse
 from collections import OrderedDict
 
@@ -152,6 +153,8 @@ def _as_odict(obj, memo):
             return force_real_str(obj)
         elif isinstance(obj, collections_abc.Iterable) and not isinstance(obj, collections_abc.Iterator):
             return type(obj)(SwaggerDict._as_odict(elem, memo) for elem in obj)
+        elif isinstance(obj, enum.Enum):
+            return obj.value
 
         return obj
 

tests/test_reference_schema.py

@@ -1,4 +1,5 @@
 import json
+import enum
 from collections import OrderedDict
 
 from drf_yasg import openapi
@@ -10,6 +11,11 @@ def test_reference_schema(swagger_dict, reference_schema, compare_schemas):
     compare_schemas(swagger_dict, reference_schema)
 
 
+class VerisonEnum(enum.Enum):
+    V1 = 'v1'
+    V2 = 'v2'
+
+
 class NoOpFieldInspector(FieldInspector):
     pass
 
@@ -38,8 +44,8 @@ def set_inspectors(inspectors, setting_name):
     set_inspectors([NoOpPaginatorInspector], 'DEFAULT_PAGINATOR_INSPECTORS')
 
     generator = OpenAPISchemaGenerator(
-        info=openapi.Info(title="Test generator", default_version="v1"),
-        version="v2",
+        info=openapi.Info(title="Test generator", default_version=VerisonEnum.V1),
+        version=VerisonEnum.V2,
     )
     swagger = generator.get_schema(mock_schema_request, True)
 

tests/test_schema_generator.py

@@ -51,14 +51,6 @@ def test_yaml_codec_roundtrip(codec_yaml, swagger, validate_schema):
     validate_schema(yaml_sane_load(yaml_bytes.decode('utf-8')))
 
 
-def test_yaml_render_enums():
-    class TestEnum(enum.Enum):
-        ONE = 'one'
-        TWO = 'two'
-
-    assert yaml_sane_dump({TestEnum.ONE: TestEnum.TWO}, False) == f"{TestEnum.ONE.value}: {TestEnum.TWO.value}\n"
-
-
 def test_yaml_and_json_match(codec_yaml, codec_json, swagger):
     yaml_schema = yaml_sane_load(codec_yaml.encode(swagger).decode('utf-8'))
     json_schema = json.loads(codec_json.encode(swagger).decode('utf-8'), object_pairs_hook=OrderedDict)