use std::{collections::BTreeSet, rc::Rc};
use anyhow::Context;
use itertools::Itertools;
use oas3::spec::{ObjectSchema, Schema, SchemaType};
use super::{
ConversionOutput,
methods::MethodGenerator,
relaxed_enum::RelaxedEnumBuilder,
union_types::{CollisionStrategy, UnionVariantSpec, variants_to_cache_key},
value_enums::ValueEnumBuilder,
variants::VariantBuilder,
};
use crate::{
generator::{
ast::{Documentation, EnumVariantToken, RustPrimitive, RustType},
converter::{ConverterContext, discriminator::DiscriminatorConverter},
naming::{identifiers::ensure_unique, inference::strip_common_affixes},
},
utils::{SchemaExt, SchemaRefName, SchemaResolveExt},
};
#[derive(Clone, Debug)]
pub(crate) struct EnumConverter {
context: Rc<ConverterContext>,
value_enum_builder: ValueEnumBuilder,
}
impl EnumConverter {
pub(crate) fn new(context: Rc<ConverterContext>) -> Self {
let case_insensitive = context.config().case_insensitive_enums();
let sort_variants = context.config().sort_enum_variants();
Self {
context,
value_enum_builder: ValueEnumBuilder::new(case_insensitive, sort_variants),
}
}
pub(crate) fn convert_value_enum(&self, name: &str, schema: &ObjectSchema) -> RustType {
let strategy = if self.context.config().preserve_case_variants() {
CollisionStrategy::Preserve
} else {
CollisionStrategy::Deduplicate
};
let variants = schema.extract_enum_entries(self.context.graph().spec());
let scalar_repr = Self::scalar_repr(schema);
self.value_enum_builder.build_enum_from_variants(
name,
variants,
strategy,
Documentation::from_optional(schema.description.as_ref()),
scalar_repr,
)
}
fn scalar_repr(schema: &ObjectSchema) -> Option<RustPrimitive> {
match schema.single_type() {
Some(SchemaType::Integer) => Some(RustPrimitive::with_format_override(
RustPrimitive::I64,
schema.format.as_deref(),
)),
Some(SchemaType::Number) => Some(RustPrimitive::with_format_override(
RustPrimitive::F64,
schema.format.as_deref(),
)),
_ => None,
}
}
}
#[derive(Clone, Debug)]
pub(crate) struct UnionConverter {
context: Rc<ConverterContext>,
variant_builder: VariantBuilder,
relaxed_enum_builder: RelaxedEnumBuilder,
method_generator: MethodGenerator,
discriminator_converter: DiscriminatorConverter,
}
impl UnionConverter {
pub(crate) fn new(context: Rc<ConverterContext>) -> Self {
let variant_builder = VariantBuilder::new(context.clone());
let relaxed_enum_builder = RelaxedEnumBuilder::new(context.clone());
let method_generator = MethodGenerator::new(context.clone());
let discriminator_converter = DiscriminatorConverter::new(context.clone());
Self {
context,
variant_builder,
relaxed_enum_builder,
method_generator,
discriminator_converter,
}
}
pub(crate) fn convert_union(&self, name: &str, schema: &ObjectSchema) -> anyhow::Result<ConversionOutput<RustType>> {
if !schema.any_of.is_empty()
&& let Some(output) = self.relaxed_enum_builder.try_build_relaxed_enum(name, schema)
{
return Ok(output);
}
let output = self.collect_union_variants(name, schema)?;
let should_register_enum = !schema.enum_values.is_empty() || schema.has_relaxed_anyof_enum();
if should_register_enum {
let variants = schema.extract_enum_entries(self.context.graph().spec());
if !variants.is_empty()
&& let RustType::Enum(e) = &output.result
{
let cache_key = variants_to_cache_key(&variants);
self
.context
.cache
.borrow_mut()
.register_enum(cache_key, e.name.to_string());
}
}
Ok(output)
}
fn collect_union_variants(&self, name: &str, schema: &ObjectSchema) -> anyhow::Result<ConversionOutput<RustType>> {
let variants_src = if schema.one_of.is_empty() {
&schema.any_of
} else {
&schema.one_of
};
let variant_specs = self.collect_union_variant_specs(variants_src)?;
let (mut variants, inline_types) = itertools::process_results(
variant_specs.into_iter().map(|spec| {
let output = self.variant_builder.build_variant(name, &spec)?;
anyhow::Ok((output.result, output.inline_types))
}),
|iter| iter.unzip::<_, _, Vec<_>, Vec<_>>(),
)?;
let inline_types = inline_types.into_iter().flatten().collect::<Vec<_>>();
variants = strip_common_affixes(variants);
if self.context.config().sort_enum_variants() {
variants = variants
.into_iter()
.sorted_by(|a, b| a.name.as_str().cmp(b.name.as_str()))
.collect();
}
let methods = if self.context.config().no_helpers() {
vec![]
} else {
self.method_generator.build_constructors(&variants, &inline_types, name)
};
let main_enum = self
.discriminator_converter
.try_upgrade_to_discriminated(name, schema, &variants, methods.clone())
.unwrap_or_else(|| {
RustType::untagged_enum()
.name(name)
.docs(Documentation::from_optional(schema.description.as_ref()))
.variants(variants)
.methods(methods)
.call()
});
Ok(ConversionOutput::with_inline_types(main_enum, inline_types))
}
fn collect_union_variant_specs(&self, variants_src: &[Schema]) -> anyhow::Result<Vec<UnionVariantSpec>> {
let mut specs = vec![];
let mut seen_names = BTreeSet::new();
for (i, variant_ref) in variants_src.iter().enumerate() {
let resolved = variant_ref
.resolve_object(self.context.graph().spec())
.context(format!("Schema resolution failed for union variant {i}"))?;
if resolved.is_null() {
continue;
}
let ref_name = variant_ref.schema_ref_name().or_else(|| {
if resolved.all_of.len() == 1 {
resolved.all_of[0].schema_ref_name()
} else {
None
}
});
let base_name = resolved.infer_union_variant_label(ref_name.as_deref(), i);
let variant_name = ensure_unique(&base_name, &seen_names);
seen_names.insert(variant_name.clone());
specs.push(UnionVariantSpec {
variant_name: EnumVariantToken::new(variant_name),
resolved_schema: resolved,
ref_name,
});
}
Ok(specs)
}
}