211 lines
9.9 KiB
Python
211 lines
9.9 KiB
Python
import re
|
|
from collections import defaultdict
|
|
|
|
from inflection import camelize
|
|
from rest_framework.settings import api_settings
|
|
|
|
from drf_spectacular.drainage import warn
|
|
from drf_spectacular.plumbing import (
|
|
ResolvedComponent, list_hash, load_enum_name_overrides, safe_ref,
|
|
)
|
|
from drf_spectacular.settings import spectacular_settings
|
|
|
|
|
|
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
|
|
every occurrence. only takes effect when replacement is guaranteed to be correct.
|
|
"""
|
|
|
|
def iter_prop_containers(schema, component_name=None):
|
|
if not component_name:
|
|
for component_name, schema in schema.items():
|
|
if spectacular_settings.COMPONENT_SPLIT_PATCH:
|
|
component_name = re.sub('^Patched(.+)', r'\1', component_name)
|
|
if spectacular_settings.COMPONENT_SPLIT_REQUEST:
|
|
component_name = re.sub('(.+)Request$', r'\1', component_name)
|
|
yield from iter_prop_containers(schema, component_name)
|
|
elif isinstance(schema, list):
|
|
for item in schema:
|
|
yield from iter_prop_containers(item, component_name)
|
|
elif isinstance(schema, dict):
|
|
if schema.get('properties'):
|
|
yield component_name, schema['properties']
|
|
yield from iter_prop_containers(schema.get('oneOf', []), component_name)
|
|
yield from iter_prop_containers(schema.get('allOf', []), component_name)
|
|
yield from iter_prop_containers(schema.get('anyOf', []), component_name)
|
|
|
|
def create_enum_component(name, schema):
|
|
component = ResolvedComponent(
|
|
name=name,
|
|
type=ResolvedComponent.SCHEMA,
|
|
schema=schema,
|
|
object=name,
|
|
)
|
|
generator.registry.register_on_missing(component)
|
|
return component
|
|
|
|
def extract_hash(schema):
|
|
if 'x-spec-enum-id' in schema:
|
|
# try to use the injected enum hash first as it generated from (name, value) tuples,
|
|
# which prevents collisions on choice sets only differing in labels not values.
|
|
return schema['x-spec-enum-id']
|
|
else:
|
|
# fall back to actual list hashing when we encounter enums not generated by us.
|
|
# remove blank/null entry for hashing. will be reconstructed in the last step
|
|
return list_hash([(i, i) for i in schema['enum'] if i not in ('', None)])
|
|
|
|
schemas = result.get('components', {}).get('schemas', {})
|
|
|
|
overrides = load_enum_name_overrides()
|
|
|
|
prop_hash_mapping = defaultdict(set)
|
|
hash_name_mapping = defaultdict(set)
|
|
# collect all enums, their names and choice sets
|
|
for component_name, props in iter_prop_containers(schemas):
|
|
for prop_name, prop_schema in props.items():
|
|
if prop_schema.get('type') == 'array':
|
|
prop_schema = prop_schema.get('items', {})
|
|
if 'enum' not in prop_schema:
|
|
continue
|
|
|
|
prop_enum_cleaned_hash = extract_hash(prop_schema)
|
|
prop_hash_mapping[prop_name].add(prop_enum_cleaned_hash)
|
|
hash_name_mapping[prop_enum_cleaned_hash].add((component_name, prop_name))
|
|
|
|
# get the suffix to be used for enums from settings
|
|
enum_suffix = spectacular_settings.ENUM_SUFFIX
|
|
|
|
# traverse all enum properties and generate a name for the choice set. naming collisions
|
|
# are resolved and a warning is emitted. giving a choice set multiple names is technically
|
|
# correct but potentially unwanted. also emit a warning there to make the user aware.
|
|
enum_name_mapping = {}
|
|
for prop_name, prop_hash_set in prop_hash_mapping.items():
|
|
for prop_hash in prop_hash_set:
|
|
if prop_hash in overrides:
|
|
enum_name = overrides[prop_hash]
|
|
elif len(prop_hash_set) == 1:
|
|
# prop_name has been used exclusively for one choice set (best case)
|
|
enum_name = f'{camelize(prop_name)}{enum_suffix}'
|
|
elif len(hash_name_mapping[prop_hash]) == 1:
|
|
# prop_name has multiple choice sets, but each one limited to one component only
|
|
component_name, _ = next(iter(hash_name_mapping[prop_hash]))
|
|
enum_name = f'{camelize(component_name)}{camelize(prop_name)}{enum_suffix}'
|
|
else:
|
|
enum_name = f'{camelize(prop_name)}{prop_hash[:3].capitalize()}{enum_suffix}'
|
|
warn(
|
|
f'enum naming encountered a non-optimally resolvable collision for fields '
|
|
f'named "{prop_name}". The same name has been used for multiple choice sets '
|
|
f'in multiple components. The collision was resolved with "{enum_name}". '
|
|
f'add an entry to ENUM_NAME_OVERRIDES to fix the naming.'
|
|
)
|
|
if enum_name_mapping.get(prop_hash, enum_name) != enum_name:
|
|
warn(
|
|
f'encountered multiple names for the same choice set ({enum_name}). This '
|
|
f'may be unwanted even though the generated schema is technically correct. '
|
|
f'Add an entry to ENUM_NAME_OVERRIDES to fix the naming.'
|
|
)
|
|
del enum_name_mapping[prop_hash]
|
|
else:
|
|
enum_name_mapping[prop_hash] = enum_name
|
|
enum_name_mapping[(prop_hash, prop_name)] = enum_name
|
|
|
|
# replace all enum occurrences with a enum schema component. cut out the
|
|
# enum, replace it with a reference and add a corresponding component.
|
|
for _, props in iter_prop_containers(schemas):
|
|
for prop_name, prop_schema in props.items():
|
|
is_array = prop_schema.get('type') == 'array'
|
|
if is_array:
|
|
prop_schema = prop_schema.get('items', {})
|
|
|
|
if 'enum' not in prop_schema:
|
|
continue
|
|
|
|
prop_enum_original_list = prop_schema['enum']
|
|
prop_schema['enum'] = [i for i in prop_schema['enum'] if i not in ['', None]]
|
|
prop_hash = extract_hash(prop_schema)
|
|
# when choice sets are reused under multiple names, the generated name cannot be
|
|
# resolved from the hash alone. fall back to prop_name and hash for resolution.
|
|
enum_name = enum_name_mapping.get(prop_hash) or enum_name_mapping[prop_hash, prop_name]
|
|
|
|
# split property into remaining property and enum component parts
|
|
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', 'x-spec-enum-id']}
|
|
|
|
# separate actual description from name-value tuples
|
|
if spectacular_settings.ENUM_GENERATE_CHOICE_DESCRIPTION:
|
|
if prop_schema.get('description', '').startswith('*'):
|
|
enum_schema['description'] = prop_schema.pop('description')
|
|
elif '\n\n*' in prop_schema.get('description', ''):
|
|
_, _, post = prop_schema['description'].partition('\n\n*')
|
|
enum_schema['description'] = '*' + post
|
|
|
|
components = [
|
|
create_enum_component(enum_name, schema=enum_schema)
|
|
]
|
|
if spectacular_settings.ENUM_ADD_EXPLICIT_BLANK_NULL_CHOICE:
|
|
if '' in prop_enum_original_list:
|
|
components.append(create_enum_component(f'Blank{enum_suffix}', schema={'enum': ['']}))
|
|
if None in prop_enum_original_list:
|
|
if spectacular_settings.OAS_VERSION.startswith('3.1'):
|
|
components.append(create_enum_component(f'Null{enum_suffix}', schema={'type': 'null'}))
|
|
else:
|
|
components.append(create_enum_component(f'Null{enum_suffix}', schema={'enum': [None]}))
|
|
|
|
# undo OAS 3.1 type list NULL construction as we cover this in a separate component already
|
|
if spectacular_settings.OAS_VERSION.startswith('3.1') and isinstance(enum_schema['type'], list):
|
|
enum_schema['type'] = [t for t in enum_schema['type'] if t != 'null'][0]
|
|
|
|
if len(components) == 1:
|
|
prop_schema.update(components[0].ref)
|
|
else:
|
|
prop_schema.update({'oneOf': [c.ref for c in components]})
|
|
|
|
if is_array:
|
|
props[prop_name]['items'] = safe_ref(prop_schema)
|
|
else:
|
|
props[prop_name] = safe_ref(prop_schema)
|
|
|
|
# sort again with additional components
|
|
result['components'] = generator.registry.build(spectacular_settings.APPEND_COMPONENTS)
|
|
|
|
# remove remaining ids that were not part of this hook (operation parameters mainly)
|
|
postprocess_schema_enum_id_removal(result, generator)
|
|
|
|
return result
|
|
|
|
|
|
def postprocess_schema_enum_id_removal(result, generator, **kwargs):
|
|
"""
|
|
Iterative modifying approach to scanning the whole schema and removing the
|
|
temporary helper ids that allowed us to distinguish similar enums.
|
|
"""
|
|
def clean(sub_result):
|
|
if isinstance(sub_result, dict):
|
|
for key in list(sub_result):
|
|
if key == 'x-spec-enum-id':
|
|
del sub_result['x-spec-enum-id']
|
|
else:
|
|
clean(sub_result[key])
|
|
elif isinstance(sub_result, (list, tuple)):
|
|
for item in sub_result:
|
|
clean(item)
|
|
|
|
clean(result)
|
|
|
|
return result
|
|
|
|
|
|
def preprocess_exclude_path_format(endpoints, **kwargs):
|
|
"""
|
|
preprocessing hook that filters out {format} suffixed paths, in case
|
|
format_suffix_patterns is used and {format} path params are unwanted.
|
|
"""
|
|
format_path = f'{{{api_settings.FORMAT_SUFFIX_KWARG}}}'
|
|
return [
|
|
(path, path_regex, method, callback)
|
|
for path, path_regex, method, callback in endpoints
|
|
if not (path.endswith(format_path) or path.endswith(format_path + '/'))
|
|
]
|