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 + '/')) ]