import dataclasses
from typing import TYPE_CHECKING, Any, Literal, get_args, get_origin, overload
from dmr.exceptions import UnsolvableAnnotationsError
from dmr.openapi.mappers.example import generate_example
from dmr.openapi.mappers.schema_loader import load_schema
from dmr.openapi.objects import Reference, Schema
if TYPE_CHECKING:
from dmr.openapi.core.context import OpenAPIContext
from dmr.serializer import BaseSerializer
[docs]
@dataclasses.dataclass(frozen=True, slots=True)
class SchemaGenerator:
"""Generate OpenAPI schemas from different type annotations."""
# Instance API:
_context: 'OpenAPIContext'
@overload
def __call__(
self,
annotation: Any,
serializer: type['BaseSerializer'],
*,
used_for_response: bool = False,
skip_registration: Literal[True],
) -> Schema: ...
@overload
def __call__(
self,
annotation: Any,
serializer: type['BaseSerializer'],
*,
used_for_response: bool = False,
skip_registration: bool = False,
) -> Schema | Reference: ...
[docs]
def __call__(
self,
annotation: Any,
serializer: type['BaseSerializer'],
*,
used_for_response: bool = False,
skip_registration: bool = False,
) -> Schema | Reference:
"""
Get schema for an annotation.
Here's the algorithm we use:
1. First, we try to find manually defined overrides for the annotation
2. If nothing is found, we try to find any existing schema references
3. Next, we try to get a model schema from a serializer.
If it exists, we create an internal reference and return it.
The next time it will be returned as a reference, cached.
4. If nothing worked, we raise an error
Raises:
UnsolvableAnnotationsError: when we can't generate
an OpenAPI schema from an existing annotation.
"""
explicit_override = self._resolve_schema_override(
annotation,
serializer,
used_for_response=used_for_response,
skip_registration=skip_registration,
)
if explicit_override:
return explicit_override
existing_reference = self._context.registries.schema.get_reference(
(
serializer.schema_generator.schema_name(annotation)
or getattr(annotation, '__qualname__', None)
),
annotation,
)
if existing_reference is not None:
return existing_reference
try:
schemas = serializer.schema_generator.get_schema(
annotation,
ref_template=self._context.registries.schema.schema_prefix,
used_for_response=used_for_response,
)
except Exception as exc:
raise UnsolvableAnnotationsError(
f'Cannot generate OpenAPI schema from {annotation}, '
'consider registering it as described in your serializer',
) from exc
return self._maybe_generate_reference(
annotation,
*schemas,
serializer,
skip_registration=skip_registration,
)
def _resolve_schema_override(
self,
annotation: Any,
serializer: type['BaseSerializer'],
*,
used_for_response: bool = False,
skip_registration: bool = False,
) -> Schema | Reference | None:
origin = get_origin(annotation) or annotation
type_args = get_args(annotation)
registry = self._context.registries.schema
schema = registry.overrides.get(origin)
if callable(schema):
return schema(
annotation,
origin,
type_args,
used_for_response=used_for_response,
skip_registration=skip_registration,
)
if schema is not None:
return schema
return None
def _maybe_generate_reference(
self,
annotation: Any,
schema: dict[str, Any],
components: dict[str, Any],
serializer: type['BaseSerializer'],
*,
skip_registration: bool = False,
) -> Schema | Reference:
if not skip_registration:
for component_name, component in components.items():
self._context.registries.schema.register(
schema_name=component_name,
schema=load_schema(component),
)
reference = schema.get('$ref')
if reference:
# We got a reference back, return it. It is registered already.
reference = Reference(
ref=reference,
summary=schema.get('summary'),
description=schema.get('description'),
)
if not skip_registration:
# If we got a reference from the start,
# it might still miss the examples:
self._maybe_generate_example(reference, annotation, serializer)
# When skip registration is requested, we need
# real schemas back, not references,
# because there's no registered schema under the reference.
return (
self._context.registries.schema.maybe_resolve_reference(
reference,
resolution_context=_build_resolution_context(components),
)
if skip_registration
else reference
)
# Register the final schema:
schema_obj = load_schema(
schema,
should_generate_example=True,
annotation=annotation,
serializer=serializer,
)
if not skip_registration and schema_obj.title:
return self._context.registries.schema.register(
schema_name=schema_obj.title,
schema=schema_obj,
annotation=annotation,
)
return schema_obj
def _maybe_generate_example(
self,
reference: Reference,
annotation: Any,
serializer: type['BaseSerializer'],
) -> None:
schema = self._context.registries.schema.maybe_resolve_reference(
reference,
)
if not schema.example and not schema.examples: # pragma: no branch
schema.example = generate_example(annotation, serializer)
def _build_resolution_context(
components: dict[str, Any],
) -> dict[str, Schema]:
return {
component_name: load_schema(component)
for component_name, component in components.items()
}
django-modern-rest