bugfix component sorting to include enums. #60

tfranzel · May 21, 2020 · 749dfb6 · 749dfb6
1 parent 5d15ef1
commit 749dfb6
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 15 deletions.
diff --git a/drf_spectacular/plumbing.py b/drf_spectacular/plumbing.py
@@ -317,7 +317,7 @@ def __init__(self):
         self._components = {}
 
     def register(self, component: ResolvedComponent):
-        if self._components.get(component.key, None):
+        if component.key in self._components:
             warn(
                 f'trying to re-register a {component.type} component with name '
                 f'{self._components[component.key].name}. this might lead to '
@@ -408,7 +408,7 @@ def get_match(cls, target) -> Optional[T]:
         return None
 
 
-def postprocess_schema_enums(result, **kwargs):
+def postprocess_schema_enums(result, generator, **kwargs):
     """
     simple replacement of Enum/Choices that globally share the same name and have
     the same choices. Aids client generation to not generate a separate enum for
@@ -424,7 +424,6 @@ def iter_prop_containers(schema):
             yield from iter_prop_containers(schema.get('oneOf', []))
             yield from iter_prop_containers(schema.get('allOf', []))
 
-    # schemas will be modified in-place and same result object is returned
     schemas = result.get('components', {}).get('schemas', {})
 
     hash_mapping = defaultdict(set)
@@ -452,11 +451,20 @@ def iter_prop_containers(schema):
             enum_name = f'{inflection.camelize(prop_name)}Enum'
             enum_schema = {k: v for k, v in prop_schema.items() if k in ['type', 'enum']}
             prop_schema = {k: v for k, v in prop_schema.items() if k not in ['type', 'enum']}
-            prop_schema['$ref'] = f'#/components/schemas/{enum_name}'
 
-            schemas[enum_name] = enum_schema
+            component = ResolvedComponent(
+                name=enum_name,
+                type=ResolvedComponent.SCHEMA,
+                schema=enum_schema,
+                object=enum_name,
+            )
+            if component not in generator.registry:
+                generator.registry.register(component)
+            prop_schema.update(component.ref)
             props[prop_name] = safe_ref(prop_schema)
 
+    # sort again with additional components
+    result['components'] = generator.registry.build(spectacular_settings.APPEND_COMPONENTS)
     return result
 
 

diff --git a/tests/test_basic.yml b/tests/test_basic.yml
@@ -218,6 +218,11 @@ components:
       - songs
       - title
       - year
+    GenreEnum:
+      enum:
+      - POP
+      - ROCK
+      type: string
     PatchedAlbum:
       type: object
       properties:
@@ -264,11 +269,6 @@ components:
       - length
       - title
       - top10
-    GenreEnum:
-      enum:
-      - POP
-      - ROCK
-      type: string
   securitySchemes:
     tokenAuth:
       type: http

diff --git a/tests/test_extend_schema.yml b/tests/test_extend_schema.yml
@@ -294,6 +294,11 @@ components:
       - field_j
       - field_k
       - field_l
+    FieldLEnum:
+      enum:
+      - a
+      - b
+      type: string
     Gamma:
       type: object
       properties:
@@ -339,11 +344,6 @@ components:
           - a
       required:
       - stars
-    FieldLEnum:
-      enum:
-      - a
-      - b
-      type: string
   securitySchemes:
     basicAuth:
       type: http