use super::errors::is_dataclass_backed_config;
use crate::codegen::generators::{
PYO3_DTO_COERCE_HELPER, coercible_payload, collect_variant_constructors, data_enum_needs_dto_coercion,
enum_has_data_variants, pyo3_wire_schema_const_name,
};
use crate::codegen::naming::wire_field_name;
use crate::codegen::shared::binding_fields;
use crate::core::config::ResolvedCrateConfig;
use crate::core::ir::{ApiSurface, TypeDef};
use ahash::{AHashMap, AHashSet};
pub(in crate::backends::pyo3) fn coercible_dto_names<'a>(
api: &'a ApiSurface,
config: &ResolvedCrateConfig,
) -> AHashSet<&'a str> {
let output_style = config.dto.python_output_style();
let reexported: AHashSet<&str> = config
.python
.as_ref()
.map(|p| p.reexported_types.iter().map(String::as_str).collect())
.unwrap_or_default();
api.types
.iter()
.filter(|t| is_dataclass_backed_config(t, output_style, &reexported))
.map(|t| t.name.as_str())
.collect()
}
pub(super) fn emit_dto_coercion_section(api: &ApiSurface, has_serde: bool, coercible: &AHashSet<&str>) -> String {
let needed = has_serde && api.enums.iter().any(|e| data_enum_needs_dto_coercion(e, coercible));
if !needed {
return String::new();
}
let schema_consts = gen_wire_schema_consts(api, coercible);
if schema_consts.is_empty() {
PYO3_DTO_COERCE_HELPER.to_string()
} else {
format!("{PYO3_DTO_COERCE_HELPER}\n\n{schema_consts}")
}
}
struct AliasEntry {
rust: String,
wire: String,
kind: &'static str,
nested: String,
}
pub(super) fn gen_wire_schema_consts(api: &ApiSurface, coercible_dto_names: &AHashSet<&str>) -> String {
if coercible_dto_names.is_empty() {
return String::new();
}
let types: AHashMap<&str, &TypeDef> = api.types.iter().map(|t| (t.name.as_str(), t)).collect();
let mut seeds: Vec<String> = Vec::new();
for e in &api.enums {
if !enum_has_data_variants(e) {
continue;
}
for ctor in collect_variant_constructors(e) {
for p in &ctor.params {
if let Some((dto, _)) = coercible_payload(&p.ty, coercible_dto_names) {
if !seeds.iter().any(|s| s == dto) {
seeds.push(dto.to_string());
}
}
}
}
}
if seeds.is_empty() {
return String::new();
}
let mut built: Vec<(String, Vec<AliasEntry>)> = Vec::new();
let mut done: AHashSet<String> = AHashSet::new();
for seed in &seeds {
build_type(
seed,
&types,
coercible_dto_names,
&mut Vec::new(),
&mut done,
&mut built,
);
}
built.sort_by(|a, b| a.0.cmp(&b.0));
let mut out = String::new();
for (const_name, entries) in &built {
let rendered_entries: Vec<minijinja::Value> = entries
.iter()
.map(|e| {
minijinja::context! {
rust => &e.rust,
wire => &e.wire,
kind => e.kind,
nested => &e.nested,
}
})
.collect();
out.push_str(&crate::backends::pyo3::template_env::render(
"pyo3_wire_schema_const.jinja",
minijinja::context! {
const_name => const_name,
entries => rendered_entries,
},
));
out.push('\n');
}
out
}
fn build_type(
type_name: &str,
types: &AHashMap<&str, &TypeDef>,
coercible: &AHashSet<&str>,
path: &mut Vec<String>,
done: &mut AHashSet<String>,
built: &mut Vec<(String, Vec<AliasEntry>)>,
) {
if done.contains(type_name) {
return;
}
let Some(typ) = types.get(type_name) else {
done.insert(type_name.to_string());
built.push((pyo3_wire_schema_const_name(type_name), Vec::new()));
return;
};
done.insert(type_name.to_string());
path.push(type_name.to_string());
let rename_all = typ.serde_rename_all.as_deref();
let mut entries: Vec<AliasEntry> = Vec::new();
for field in binding_fields(&typ.fields) {
let wire = wire_field_name(&field.name, field.serde_rename.as_deref(), rename_all);
match coercible_payload(&field.ty, coercible) {
Some((dto_name, shape)) => {
let nested = if path.iter().any(|p| p == dto_name) {
"&[]".to_string()
} else {
build_type(dto_name, types, coercible, path, done, built);
pyo3_wire_schema_const_name(dto_name)
};
entries.push(AliasEntry {
rust: field.name.clone(),
wire,
kind: shape.wire_kind(),
nested,
});
}
None => {
if wire != field.name {
entries.push(AliasEntry {
rust: field.name.clone(),
wire,
kind: "Leaf",
nested: "&[]".to_string(),
});
}
}
}
}
path.pop();
built.push((pyo3_wire_schema_const_name(type_name), entries));
}
#[cfg(test)]
mod tests;