Merge pull request #135 from jriggins/enum

tomchristie · web-flow · commit 5ae907e61113 · 2017-05-23T12:07:15.000-07:00
Fix the encoding/decoding of Enum types.
diff --git a/coreapi/codecs/corejson.py b/coreapi/codecs/corejson.py
@@ -33,19 +33,27 @@
 
 def encode_schema_to_corejson(schema):
     type_id = SCHEMA_CLASS_TO_TYPE_ID.get(schema.__class__, 'anything')
-    return {
+    retval = {
         '_type': type_id,
         'title': schema.title,
         'description': schema.description
     }
+    if isinstance(schema, coreschema.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')
+
+    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)
+    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
@@ -4,6 +4,7 @@
 from coreapi.document import Document, Link, Error, Field
 from coreapi.exceptions import ParseError, NoCodecAvailable
 from coreapi.utils import negotiate_decoder, negotiate_encoder
+from coreschema import Enum, String
 import pytest
 
 
@@ -21,7 +22,13 @@ def doc():
             'integer': 123,
             'dict': {'key': 'value'},
             'list': [1, 2, 3],
-            'link': Link(url='http://example.org/', fields=[Field(name='example')]),
+            'link': Link(
+                url='http://example.org/',
+                fields=[
+                    Field(name='noschema'),
+                    Field(name='string_example', schema=String()),
+                    Field(name='enum_example', schema=Enum(['a', 'b', 'c'])),
+                ]),
             'nested': {'child': Link(url='http://example.org/123')},
             '_type': 'needs escaping'
         })
@@ -40,7 +47,26 @@ def test_document_to_primitive(doc):
         'integer': 123,
         'dict': {'key': 'value'},
         'list': [1, 2, 3],
-        'link': {'_type': 'link', 'fields': [{'name': 'example'}]},
+        'link': {'_type': 'link', 'fields': [
+            {'name': 'noschema'},
+            {
+                'name': 'string_example',
+                'schema': {
+                    '_type': 'string',
+                    'title': '',
+                    'description': '',
+                },
+            },
+            {
+                'name': 'enum_example',
+                'schema': {
+                    '_type': 'enum',
+                    'title': '',
+                    'description': '',
+                    'enum': ['a', 'b', 'c'],
+                },
+            },
+        ]},
         'nested': {'child': {'_type': 'link', 'url': '/123'}},
         '__type': 'needs escaping'
     }
@@ -56,7 +82,30 @@ def test_primitive_to_document(doc):
         'integer': 123,
         'dict': {'key': 'value'},
         'list': [1, 2, 3],
-        'link': {'_type': 'link', 'url': 'http://example.org/', 'fields': [{'name': 'example'}]},
+        'link': {
+            '_type': 'link',
+            'url': 'http://example.org/',
+            'fields': [
+                {'name': 'noschema'},
+                {
+                    'name': 'string_example',
+                    'schema': {
+                        '_type': 'string',
+                        'title': '',
+                        'description': '',
+                    },
+                },
+                {
+                    'name': 'enum_example',
+                    'schema': {
+                        '_type': 'enum',
+                        'title': '',
+                        'description': '',
+                        'enum': ['a', 'b', 'c'],
+                    },
+                },
+            ],
+        },
         'nested': {'child': {'_type': 'link', 'url': 'http://example.org/123'}},
         '__type': 'needs escaping'
     }
