95 lines
4.4 KiB
Python
95 lines
4.4 KiB
Python
from drf_spectacular.drainage import error, warn
|
|
from drf_spectacular.extensions import OpenApiSerializerExtension
|
|
from drf_spectacular.plumbing import force_instance, is_list_serializer, is_serializer
|
|
|
|
|
|
class PolymorphicProxySerializerExtension(OpenApiSerializerExtension):
|
|
target_class = 'drf_spectacular.utils.PolymorphicProxySerializer'
|
|
priority = -1
|
|
|
|
def get_name(self):
|
|
return self.target.component_name
|
|
|
|
def map_serializer(self, auto_schema, direction):
|
|
""" custom handling for @extend_schema's injection of PolymorphicProxySerializer """
|
|
if isinstance(self.target.serializers, dict):
|
|
sub_components = self._get_explicit_sub_components(auto_schema, direction)
|
|
else:
|
|
sub_components = self._get_implicit_sub_components(auto_schema, direction)
|
|
|
|
if not self._has_discriminator():
|
|
return {'oneOf': [schema for _, schema in sub_components]}
|
|
else:
|
|
one_of_list = []
|
|
for _, schema in sub_components:
|
|
if schema not in one_of_list:
|
|
one_of_list.append(schema)
|
|
return {
|
|
'oneOf': one_of_list,
|
|
'discriminator': {
|
|
'propertyName': self.target.resource_type_field_name,
|
|
'mapping': {resource_type: schema['$ref'] for resource_type, schema in sub_components}
|
|
}
|
|
}
|
|
|
|
def _get_implicit_sub_components(self, auto_schema, direction):
|
|
sub_components = []
|
|
for sub_serializer in self.target.serializers:
|
|
sub_serializer = self._prep_serializer(sub_serializer)
|
|
(resolved_name, resolved_schema) = self._process_serializer(auto_schema, sub_serializer, direction)
|
|
if not resolved_schema:
|
|
continue
|
|
|
|
if not self._has_discriminator():
|
|
sub_components.append((None, resolved_schema))
|
|
else:
|
|
try:
|
|
discriminator_field = sub_serializer.fields[self.target.resource_type_field_name]
|
|
resource_type = discriminator_field.to_representation(None)
|
|
except: # noqa: E722
|
|
warn(
|
|
f'sub-serializer {resolved_name} of {self.target.component_name} '
|
|
f'must contain the discriminator field "{self.target.resource_type_field_name}". '
|
|
f'defaulting to sub-serializer name, but schema will likely not match the API.'
|
|
)
|
|
resource_type = resolved_name
|
|
sub_components.append((resource_type, resolved_schema))
|
|
|
|
return sub_components
|
|
|
|
def _get_explicit_sub_components(self, auto_schema, direction):
|
|
sub_components = []
|
|
for resource_type, sub_serializer in self.target.serializers.items():
|
|
sub_serializer = self._prep_serializer(sub_serializer)
|
|
(_, resolved_schema) = self._process_serializer(auto_schema, sub_serializer, direction)
|
|
if resolved_schema:
|
|
sub_components.append((resource_type, resolved_schema))
|
|
|
|
return sub_components
|
|
|
|
def _has_discriminator(self):
|
|
return self.target.resource_type_field_name is not None
|
|
|
|
def _prep_serializer(self, serializer):
|
|
serializer = force_instance(serializer)
|
|
serializer.partial = self.target.partial
|
|
return serializer
|
|
|
|
def _process_serializer(self, auto_schema, serializer, direction):
|
|
if is_list_serializer(serializer):
|
|
if self._has_discriminator() or self.target._many is not False:
|
|
warn(
|
|
"To control sub-serializer's \"many\" attribute, following usage pattern is necessary: "
|
|
"PolymorphicProxySerializer(serializers=[...], resource_type_field_name=None, "
|
|
"many=False). Ignoring serializer as it is not processable in this configuration."
|
|
)
|
|
return None, None
|
|
schema = auto_schema._unwrap_list_serializer(serializer, direction)
|
|
return None, schema
|
|
elif is_serializer(serializer):
|
|
resolved = auto_schema.resolve_serializer(serializer, direction)
|
|
return (resolved.name, resolved.ref) if resolved else (None, None)
|
|
else:
|
|
error("PolymorphicProxySerializer's serializer argument contained unknown objects.")
|
|
return None, None
|