authentik.api.v3.schema.enum

  1"""Error Response schema, from https://github.com/axnsan12/drf-yasg/issues/224"""
  2
  3import functools
  4import inspect
  5import re
  6from collections import defaultdict
  7from enum import Enum
  8
  9from django.db.models import Choices
 10from django.utils.translation import get_language
 11from drf_spectacular.drainage import error, warn
 12from drf_spectacular.hooks import postprocess_schema_enum_id_removal
 13from drf_spectacular.plumbing import (
 14    ResolvedComponent,
 15    deep_import_string,
 16    list_hash,
 17    safe_ref,
 18)
 19from drf_spectacular.settings import spectacular_settings
 20from inflection import camelize
 21from structlog.stdlib import get_logger
 22
 23LOGGER = get_logger()
 24
 25
 26# See https://github.com/tfranzel/drf-spectacular/blob/master/drf_spectacular/hooks.py
 27# and https://github.com/tfranzel/drf-spectacular/issues/520
 28def postprocess_schema_enums(result, generator, **kwargs):  # noqa: PLR0912, PLR0915
 29    """
 30    simple replacement of Enum/Choices that globally share the same name and have
 31    the same choices. Aids client generation to not generate a separate enum for
 32    every occurrence. only takes effect when replacement is guaranteed to be correct.
 33    """
 34
 35    def is_enum_prop(prop_schema):
 36        return (
 37            "enum" in prop_schema
 38            or prop_schema.get("type") == "array"
 39            and "enum" in prop_schema.get("items", {})
 40        )
 41
 42    def iter_field_schemas():
 43        def iter_prop_containers(schema, component_name=None):
 44            if not component_name:
 45                for _component_name, _schema in schema.items():
 46                    if spectacular_settings.COMPONENT_SPLIT_PATCH:
 47                        _component_name = re.sub("^Patched(.+)", r"\1", _component_name)
 48                    if spectacular_settings.COMPONENT_SPLIT_REQUEST:
 49                        _component_name = re.sub("(.+)Request$", r"\1", _component_name)
 50                    yield from iter_prop_containers(_schema, _component_name)
 51            elif isinstance(schema, list):
 52                for item in schema:
 53                    yield from iter_prop_containers(item, component_name)
 54            elif isinstance(schema, dict):
 55                if schema.get("properties"):
 56                    yield component_name, schema["properties"]
 57                yield from iter_prop_containers(schema.get("oneOf", []), component_name)
 58                yield from iter_prop_containers(schema.get("allOf", []), component_name)
 59                yield from iter_prop_containers(schema.get("anyOf", []), component_name)
 60
 61        def iter_path_parameters():
 62            for path in result.get("paths", {}).values():
 63                for operation in path.values():
 64                    for parameter in operation.get("parameters", []):
 65                        parameter_schema = parameter.get("schema", {})
 66                        if is_enum_prop(parameter_schema):
 67                            # Move description into enum schema
 68                            if "description" in parameter:
 69                                parameter_schema["description"] = parameter.pop("description")
 70                        if "name" not in parameter:
 71                            continue
 72                        yield "", {parameter["name"]: parameter_schema}
 73
 74        component_schemas = result.get("components", {}).get("schemas", {})
 75
 76        yield from iter_prop_containers(component_schemas)
 77        yield from iter_path_parameters()
 78
 79    def create_enum_component(name, schema):
 80        component = ResolvedComponent(
 81            name=name,
 82            type=ResolvedComponent.SCHEMA,
 83            schema=schema,
 84            object=name,
 85        )
 86        generator.registry.register_on_missing(component)
 87        return component
 88
 89    def extract_hash(schema):
 90        if "x-spec-enum-id" in schema:
 91            # try to use the injected enum hash first as it generated from (name, value) tuples,
 92            # which prevents collisions on choice sets only differing in labels not values.
 93            return schema["x-spec-enum-id"]
 94        else:
 95            # fall back to actual list hashing when we encounter enums not generated by us.
 96            # remove blank/null entry for hashing. will be reconstructed in the last step
 97            return list_hash([(i, i) for i in schema["enum"] if i not in ("", None)])
 98
 99    overrides = load_enum_name_overrides()
100
101    prop_hash_mapping = defaultdict(set)
102    hash_name_mapping = defaultdict(set)
103    # collect all enums, their names and choice sets
104    for component_name, props in iter_field_schemas():
105        for prop_name, prop_schema in props.items():
106            _prop_schema = prop_schema
107            if prop_schema.get("type") == "array":
108                _prop_schema = prop_schema.get("items", {})
109            if "enum" not in _prop_schema:
110                continue
111
112            prop_enum_cleaned_hash = extract_hash(_prop_schema)
113            prop_hash_mapping[prop_name].add(prop_enum_cleaned_hash)
114            hash_name_mapping[prop_enum_cleaned_hash].add((component_name, prop_name))
115
116    # get the suffix to be used for enums from settings
117    enum_suffix = spectacular_settings.ENUM_SUFFIX
118
119    # traverse all enum properties and generate a name for the choice set. naming collisions
120    # are resolved and a warning is emitted. giving a choice set multiple names is technically
121    # correct but potentially unwanted. also emit a warning there to make the user aware.
122    enum_name_mapping = {}
123    for prop_name, prop_hash_set in prop_hash_mapping.items():
124        for prop_hash in prop_hash_set:
125            if prop_hash in overrides:
126                enum_name = overrides[prop_hash]
127            elif len(prop_hash_set) == 1:
128                # prop_name has been used exclusively for one choice set (best case)
129                enum_name = f"{camelize(prop_name)}{enum_suffix}"
130            elif len(hash_name_mapping[prop_hash]) == 1:
131                # prop_name has multiple choice sets, but each one limited to one component only
132                component_name, _ = next(iter(hash_name_mapping[prop_hash]))
133                enum_name = f"{camelize(component_name)}{camelize(prop_name)}{enum_suffix}"
134            else:
135                enum_name = f"{camelize(prop_name)}{prop_hash[:3].capitalize()}{enum_suffix}"
136                warn(
137                    f"enum naming encountered a non-optimally resolvable collision for fields "
138                    f'named "{prop_name}". The same name has been used for multiple choice sets '
139                    f'in multiple components. The collision was resolved with "{enum_name}". '
140                    f"add an entry to ENUM_NAME_OVERRIDES to fix the naming."
141                )
142            if enum_name_mapping.get(prop_hash, enum_name) != enum_name:
143                warn(
144                    f"encountered multiple names for the same choice set ({enum_name}). This "
145                    f"may be unwanted even though the generated schema is technically correct. "
146                    f"Add an entry to ENUM_NAME_OVERRIDES to fix the naming."
147                )
148                del enum_name_mapping[prop_hash]
149            else:
150                enum_name_mapping[prop_hash] = enum_name
151            enum_name_mapping[(prop_hash, prop_name)] = enum_name
152
153    # replace all enum occurrences with a enum schema component. cut out the
154    # enum, replace it with a reference and add a corresponding component.
155    for _, props in iter_field_schemas():
156        for prop_name, _prop_schema in props.items():
157            prop_schema = _prop_schema
158            is_array = prop_schema.get("type") == "array"
159            if is_array:
160                prop_schema = prop_schema.get("items", {})
161
162            if "enum" not in prop_schema:
163                continue
164
165            prop_enum_original_list = prop_schema["enum"]
166            prop_schema["enum"] = [i for i in prop_schema["enum"] if i not in ["", None]]
167            prop_hash = extract_hash(prop_schema)
168            # when choice sets are reused under multiple names, the generated name cannot be
169            # resolved from the hash alone. fall back to prop_name and hash for resolution.
170            enum_name = enum_name_mapping.get(prop_hash) or enum_name_mapping[prop_hash, prop_name]
171
172            # split property into remaining property and enum component parts
173            enum_schema = {k: v for k, v in prop_schema.items() if k in ["type", "enum"]}
174            prop_schema = {
175                k: v for k, v in prop_schema.items() if k not in ["type", "enum", "x-spec-enum-id"]
176            }
177
178            # separate actual description from name-value tuples
179            if spectacular_settings.ENUM_GENERATE_CHOICE_DESCRIPTION:
180                if prop_schema.get("description", "").startswith("*"):
181                    enum_schema["description"] = prop_schema.pop("description")
182                elif "\n\n*" in prop_schema.get("description", ""):
183                    _, _, post = prop_schema["description"].partition("\n\n*")
184                    enum_schema["description"] = "*" + post
185
186            components = [create_enum_component(enum_name, schema=enum_schema)]
187            if spectacular_settings.ENUM_ADD_EXPLICIT_BLANK_NULL_CHOICE:
188                if "" in prop_enum_original_list:
189                    components.append(
190                        create_enum_component(f"Blank{enum_suffix}", schema={"enum": [""]})
191                    )
192                if None in prop_enum_original_list:
193                    if spectacular_settings.OAS_VERSION.startswith("3.1"):
194                        components.append(
195                            create_enum_component(f"Null{enum_suffix}", schema={"type": "null"})
196                        )
197                    else:
198                        components.append(
199                            create_enum_component(f"Null{enum_suffix}", schema={"enum": [None]})
200                        )
201
202            # undo OAS 3.1 type list NULL construction as we cover
203            # this in a separate component already
204            if spectacular_settings.OAS_VERSION.startswith("3.1") and isinstance(
205                enum_schema["type"], list
206            ):
207                enum_schema["type"] = [t for t in enum_schema["type"] if t != "null"][0]
208
209            if len(components) == 1:
210                prop_schema.update(components[0].ref)
211            else:
212                prop_schema.update({"oneOf": [c.ref for c in components]})
213
214            patch_target = props[prop_name]  # noqa: PLR1733
215            if is_array:
216                patch_target = patch_target["items"]
217
218            # Replace existing schema information with reference
219            patch_target.clear()
220            patch_target.update(safe_ref(prop_schema))
221
222    # sort again with additional components
223    result["components"] = generator.registry.build(spectacular_settings.APPEND_COMPONENTS)
224
225    # remove remaining ids that were not part of this hook (operation parameters mainly)
226    postprocess_schema_enum_id_removal(result, generator)
227
228    return result
229
230
231# Fixed version of `load_enum_name_overrides()` with a LRU cache based on language
232# *and* enum overrides.
233# Without this, API generation breaks if there is more than 1 API present (such as in split APIs)
234# Original source: drf-spectacular/drf_spectacular/plumbing.py
235def load_enum_name_overrides():
236    cache_key = get_language() or ""
237
238    for k, v in sorted(spectacular_settings.ENUM_NAME_OVERRIDES.items()):
239        cache_key += f";{k}:{v}"
240
241    return _load_enum_name_overrides(cache_key)
242
243
244# Original source: drf-spectacular/drf_spectacular/plumbing.py
245# Only change: cache_key argument instead of language.
246@functools.lru_cache
247def _load_enum_name_overrides(cache_key):
248    overrides = {}
249    for name, _choices in spectacular_settings.ENUM_NAME_OVERRIDES.items():
250        choices = _choices
251        if isinstance(choices, str):
252            choices = deep_import_string(choices)
253        if not choices:
254            warn(
255                f"unable to load choice override for {name} from ENUM_NAME_OVERRIDES. "
256                f"please check module path string."
257            )
258            continue
259        if inspect.isclass(choices) and issubclass(choices, Choices):
260            choices = choices.choices
261        if inspect.isclass(choices) and issubclass(choices, Enum):
262            choices = [(c.value, c.name) for c in choices]
263        normalized_choices = []
264        for choice in choices:
265            # Allow None values in the simple values list case
266            if isinstance(choice, str) or choice is None:
267                # TODO warning
268                normalized_choices.append((choice, choice))  # simple choice list
269            elif isinstance(choice[1], (list, tuple)):
270                normalized_choices.extend(choice[1])  # categorized nested choices
271            else:
272                normalized_choices.append(choice)  # normal 2-tuple form
273
274        # Get all of choice values that should be used in the hash, blank and
275        # None values get excluded in the post-processing hook for enum overrides,
276        # so we do the same here to ensure the hashes match
277        hashable_values = [
278            (value, label) for value, label in normalized_choices if value not in ["", None]
279        ]
280        overrides[list_hash(hashable_values)] = name
281
282    if len(spectacular_settings.ENUM_NAME_OVERRIDES) != len(overrides):
283        error(
284            "ENUM_NAME_OVERRIDES has duplication issues. Encountered multiple names "
285            "for the same choice set. Enum naming might be unexpected."
286        )
287    return overrides