Fix the encoding/decoding of Enum types.

jriggins · jriggins · commit 13fa37f58d08 · 2017-05-23T11:59:21.000-07:00
This is an update from tbeadle/enum branch with changes that tomchriste
suggested.
diff --git a/coreapi/codecs/corejson.py b/coreapi/codecs/corejson.py
@@ -39,17 +39,21 @@ def encode_schema_to_corejson(schema):
         'description': schema.description
     }
     if isinstance(schema, coreschema.Enum):
-        retval['extra'] = {'enum': schema.enum}
+        retval['enum'] = schema.enum
     return retval
 
 
 def decode_schema_from_corejson(data):
     type_id = _get_string(data, '_type')
     title = _get_string(data, 'title')
     description = _get_string(data, 'description')
-    extra = _get_dict(data, 'extra')
+
+    kwargs = {}
+    if type_id == 'enum':
+        kwargs['enum'] = _get_list(data, 'enum')
+
     schema_cls = TYPE_ID_TO_SCHEMA_CLASS.get(type_id, coreschema.Anything)
-    return schema_cls(title=title, description=description, **extra)
+    return schema_cls(title=title, description=description, **kwargs)
 
 
 # Robust dictionary lookups, that always return an item of the correct
diff --git a/tests/test_codecs.py b/tests/test_codecs.py
@@ -63,7 +63,7 @@ def test_document_to_primitive(doc):
                     '_type': 'enum',
                     'title': '',
                     'description': '',
-                    'extra': {'enum': ['a', 'b', 'c']},
+                    'enum': ['a', 'b', 'c'],
                 },
             },
         ]},
@@ -101,7 +101,7 @@ def test_primitive_to_document(doc):
                         '_type': 'enum',
                         'title': '',
                         'description': '',
-                        'extra': {'enum': ['a', 'b', 'c']},
+                        'enum': ['a', 'b', 'c'],
                     },
                 },
             ],
