postprocessing

Author: tfranzel

drf_spectacular/plumbing.py

@@ -1,10 +1,12 @@
 import inspect
+import json
 import sys
 from abc import ABCMeta
 from collections import defaultdict
 from collections.abc import Hashable
 from typing import List, Type, Optional, TypeVar, Union, Generic
 
+import inflection
 from django import __version__ as DJANGO_VERSION
 from django.utils.module_loading import import_string
 from rest_framework import fields, serializers
@@ -285,6 +287,10 @@ def build(self, extra_components) -> dict:
         for extra_type, extra_component_dict in extra_components.items():
             for component_name, component_schema in extra_component_dict.items():
                 output[extra_type][component_name] = component_schema
+
+        if 'schemas' in output:
+            postprocess_schema_enums(output['schemas'])
+
         # sort by component type then by name
         return {
             type: {name: output[type][name] for name in output[type].keys()}
@@ -329,3 +335,37 @@ def get_match(cls, target) -> Optional[T]:
             if extension._matches(target):
                 return extension(target)
         return None
+
+
+def postprocess_schema_enums(schemas):
+    """
+    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
+    every occurrence. only takes effect when replacement is guaranteed to be correct.
+    """
+    hash_mapping = defaultdict(set)
+    # collect all enums, their names and contents
+    for schema in schemas.values():
+        for prop_name, prop_schema in schema.get('properties', {}).items():
+            if 'enum' not in prop_schema:
+                continue
+            hash_mapping[prop_name].add(
+                hash(json.dumps(prop_schema, sort_keys=True))
+            )
+    # safe replacement requires name to have only one set of enum values
+    candidate_enums = {
+        prop_name for prop_name, prop_hash_set in hash_mapping.items()
+        if len(prop_hash_set) == 1
+    }
+    # replace all valid occurrences with enum schema component
+    for schema_name in list(schemas):
+        for prop_name, prop_schema in schemas[schema_name].get('properties', {}).items():
+            if 'enum' not in prop_schema:
+                continue
+            if prop_name not in candidate_enums:
+                continue
+            enum_name = f'{inflection.camelize(prop_name)}Enum'
+            schemas[enum_name] = prop_schema
+            schemas[schema_name]['properties'][prop_name] = {
+                '$ref': f'#/components/schemas/{enum_name}'
+            }

requirements/base.txt

@@ -2,4 +2,5 @@ Django>=2.2
 djangorestframework>=3.10
 uritemplate>=3.0.0
 PyYAML>=5.1
-jsonschema>=3.2.0
+jsonschema>=3.2.0
+inflection>=0.4.0

tests/test_basic.yml

@@ -210,10 +210,7 @@ components:
           type: string
           maxLength: 100
         genre:
-          enum:
-          - POP
-          - ROCK
-          type: string
+          $ref: '#/components/schemas/GenreEnum'
         year:
           type: integer
         released:
@@ -262,11 +259,13 @@ components:
           type: string
           maxLength: 100
         genre:
-          enum:
-          - POP
-          - ROCK
-          type: string
+          $ref: '#/components/schemas/GenreEnum'
         year:
           type: integer
         released:
           type: boolean
+    GenreEnum:
+      enum:
+      - POP
+      - ROCK
+      type: string

tests/test_extend_schema.yml

@@ -286,11 +286,7 @@ components:
             $ref: '#/components/schemas/Inline'
           readOnly: true
         field_l:
-          enum:
-          - a
-          - b
-          type: string
-          readOnly: true
+          $ref: '#/components/schemas/FieldLEnum'
     Inline:
       type: object
       properties:
@@ -323,3 +319,9 @@ components:
           minLength: 3
       required:
       - stars
+    FieldLEnum:
+      enum:
+      - a
+      - b
+      type: string
+      readOnly: true

tests/test_postprocessing.py

@@ -0,0 +1,34 @@
+from rest_framework import serializers, viewsets, mixins
+from rest_framework.decorators import action
+
+from drf_spectacular.utils import extend_schema
+from tests import assert_schema, generate_schema
+
+language_choices = (
+    ('en', 'en'),
+    ('es', 'es'),
+    ('ru', 'ru'),
+    ('cn', 'cn'),
+)
+
+
+class ASerializer(serializers.Serializer):
+    language = serializers.ChoiceField(choices=language_choices)
+
+
+class BSerializer(serializers.Serializer):
+    language = serializers.ChoiceField(choices=language_choices)
+
+
+class AViewset(mixins.ListModelMixin, viewsets.GenericViewSet):
+    serializer_class = ASerializer
+
+    @extend_schema(responses=BSerializer)
+    @action(detail=False, serializer_class=BSerializer)
+    def selection(self, request):
+        pass
+
+
+def test_postprocessing(no_warnings):
+    schema = generate_schema('a', AViewset)
+    assert_schema(schema, 'tests/test_postprocessing.yml')

tests/test_postprocessing.yml

@@ -0,0 +1,72 @@
+openapi: 3.0.3
+info:
+  title: ''
+  version: 0.0.0
+paths:
+  /a/:
+    get:
+      operationId: a_list
+      description: ''
+      tags:
+      - a
+      security:
+      - cookieAuth: []
+      - basicAuth: []
+      - {}
+      responses:
+        '200':
+          content:
+            application/json:
+              schema:
+                type: array
+                items:
+                  $ref: '#/components/schemas/A'
+          description: ''
+  /a/selection/:
+    get:
+      operationId: a_selection_retrieve
+      description: ''
+      tags:
+      - a
+      security:
+      - cookieAuth: []
+      - basicAuth: []
+      - {}
+      responses:
+        '200':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/B'
+          description: ''
+components:
+  securitySchemes:
+    cookieAuth:
+      type: apiKey
+      in: cookie
+      name: Session
+    basicAuth:
+      type: http
+      scheme: basic
+  schemas:
+    A:
+      type: object
+      properties:
+        language:
+          $ref: '#/components/schemas/LanguageEnum'
+      required:
+      - language
+    B:
+      type: object
+      properties:
+        language:
+          $ref: '#/components/schemas/LanguageEnum'
+      required:
+      - language
+    LanguageEnum:
+      enum:
+      - en
+      - es
+      - ru
+      - cn
+      type: string