openapi_nexus/generators/python/httpx/
emit_models.rs

1//! Model emission for IR schemas (Python dataclasses/enums/aliases).
2//!
3//! Uses sigil-stitch high-level APIs (TypeSpec, FieldSpec, FunSpec, TypeName,
4//! sigil_quote!) for structured code generation with automatic import tracking.
5//! Each schema produces one `.py` file via `FileSpec`.
6
7use crate::codegen::traits::file_writer::FileInfo;
8pub use crate::generators::python::type_names::python_type_name;
9use crate::generators::python::type_names::python_type_name_for_schema;
10use crate::generators::request_inputs::{
11    RequestInputField, RequestInputFieldKind, RequestInputModel, RequestInputPlan,
12};
13use crate::ir::types::{
14    IrEnum, IrEnumValueType, IrIntersection, IrObject, IrPrimitive, IrProperty, IrSchema,
15    IrSchemaKind, IrSpec, IrTaggedUnion, IrTaggedVariant, IrTypeExpr, IrUnion, TaggingStyle,
16};
17use heck::{ToPascalCase, ToSnakeCase};
18use sigil_stitch::code_block::CodeBlock;
19use sigil_stitch::lang::python::Python;
20use sigil_stitch::prelude::*;
21
22/// Generate every model file from the IR.
23pub fn generate_model_files(
24    ir: &IrSpec,
25    header: &str,
26    request_inputs: &RequestInputPlan,
27) -> Result<Vec<FileInfo>, String> {
28    let mut files = Vec::new();
29    for (_name, schema) in &ir.schemas {
30        let body = emit_model_body(schema, ir).ok_or_else(|| {
31            format!(
32                "unsupported schema kind for {}: {:?}",
33                schema.name, schema.kind
34            )
35        })?;
36        let stem = schema.name.to_snake_case();
37        let filename = format!("{stem}.py");
38        let mut content = String::with_capacity(header.len() + body.len());
39        content.push_str(header);
40        content.push_str(&body);
41        files.push(FileInfo::model(filename, content));
42    }
43    for model in request_inputs.models() {
44        files.push(request_input_model_file(model, header));
45    }
46    Ok(files)
47}
48
49fn request_input_model_file(model: &RequestInputModel, header: &str) -> FileInfo {
50    let class_name = model.name.to_pascal_case();
51    let filename = format!("{}.py", model.name.to_snake_case());
52    let mut imports = std::collections::BTreeSet::new();
53    let mut needs_upload = false;
54    for field in &model.fields {
55        if field.is_upload() {
56            needs_upload = true;
57        } else {
58            collect_request_input_imports(&field.type_expr, &mut imports);
59        }
60    }
61
62    let mut content = String::new();
63    content.push_str(header);
64    content.push_str("from __future__ import annotations\n\n");
65    content.push_str("from dataclasses import dataclass\n");
66    if needs_upload {
67        content.push_str("from ..runtime import UploadFile\n");
68    }
69    for import in &imports {
70        content.push_str(import);
71        content.push('\n');
72    }
73    content.push('\n');
74    content.push_str("@dataclass\n");
75    content.push_str(&format!("class {class_name}:\n"));
76    if model.fields.is_empty() {
77        content.push_str("    pass\n");
78    } else {
79        let required = model.fields.iter().filter(|field| field.required);
80        let optional = model.fields.iter().filter(|field| !field.required);
81        for field in required.chain(optional) {
82            let field_name = python_field_name(&field.wire_name);
83            let ty = request_input_python_type(field);
84            if field.required {
85                content.push_str(&format!("    {field_name}: {ty}\n"));
86            } else {
87                content.push_str(&format!("    {field_name}: {ty} | None = None\n"));
88            }
89        }
90    }
91
92    FileInfo::model(filename, content)
93}
94
95fn request_input_python_type(field: &RequestInputField) -> String {
96    match field.kind {
97        RequestInputFieldKind::UploadFile { .. } => "UploadFile".to_string(),
98        RequestInputFieldKind::SchemaValue => python_type_str(&field.type_expr),
99    }
100}
101
102fn emit_model_body(schema: &IrSchema, ir: &IrSpec) -> Option<String> {
103    let file_spec = match &schema.kind {
104        IrSchemaKind::Object(obj) => emit_object(schema, obj, ir),
105        IrSchemaKind::Enum(en) => emit_enum(schema, en),
106        IrSchemaKind::Alias(expr) => emit_alias(schema, expr),
107        IrSchemaKind::Union(u) => emit_union(schema, u),
108        IrSchemaKind::Intersection(i) => emit_intersection(schema, i, ir),
109        IrSchemaKind::TaggedUnion(tu) => emit_tagged_union(schema, tu, ir),
110    }?;
111    file_spec.render(100).ok()
112}
113
114pub fn future_annotations_header() -> CodeBlock {
115    CodeBlock::of("from __future__ import annotations", ()).expect("static header")
116}
117
118// ---------------------------------------------------------------------------
119// Object -> @dataclass
120// ---------------------------------------------------------------------------
121
122fn emit_object(schema: &IrSchema, obj: &IrObject, ir: &IrSpec) -> Option<FileSpec> {
123    let name = schema.name.to_pascal_case();
124
125    let mut file =
126        FileSpec::builder_with("model.py", Python::new()).header(future_annotations_header());
127
128    if needs_typing_literal_in_props(&obj.properties) {
129        file = file.add_import(ImportSpec::named("typing", "Literal"));
130    }
131
132    // Import tagged-union helper functions for fields referencing TaggedUnion schemas
133    for (_, prop) in &obj.properties {
134        for named_ref in collect_tagged_union_refs(&prop.type_expr, ir) {
135            let snake = named_ref.to_snake_case();
136            let module = format!(".{snake}");
137            file = file.add_import(ImportSpec::named(&module, &format!("{snake}_from_dict")));
138            file = file.add_import(ImportSpec::named(&module, &format!("{snake}_to_dict")));
139        }
140    }
141
142    let dataclass_tn = TypeName::importable("dataclasses", "dataclass");
143    let mut cls = TypeSpec::builder(&name, TypeKind::Class)
144        .annotate(AnnotationSpec::importable(dataclass_tn));
145
146    if let Some(doc) = &schema.description {
147        cls = cls.doc(&format!("{}.", escape_docstring(doc)));
148    }
149
150    let mut required: Vec<(&String, &IrProperty)> = Vec::new();
151    let mut optional: Vec<(&String, &IrProperty)> = Vec::new();
152    for (json_name, prop) in &obj.properties {
153        if prop.required && !prop.nullable {
154            required.push((json_name, prop));
155        } else {
156            optional.push((json_name, prop));
157        }
158    }
159
160    let all_fields: Vec<(&String, &IrProperty)> =
161        required.iter().chain(optional.iter()).copied().collect();
162
163    if all_fields.is_empty() {
164        cls = cls.extra_member(CodeBlock::of("pass", ()).expect("pass"));
165    } else {
166        for (_json_name, prop) in &required {
167            let field_name = python_field_name(&prop.name);
168            let type_name = python_type_name_for_schema(&prop.type_expr, &schema.name);
169            cls = cls.add_field(
170                FieldSpec::builder(&field_name, type_name)
171                    .build()
172                    .expect("required field"),
173            );
174        }
175        for (_json_name, prop) in &optional {
176            let field_name = python_field_name(&prop.name);
177            let type_name = python_type_name_for_schema(&prop.type_expr, &schema.name);
178            cls = cls.add_field(
179                FieldSpec::builder(&field_name, TypeName::optional(type_name))
180                    .initializer(CodeBlock::of("None", ()).expect("None init"))
181                    .build()
182                    .expect("optional field"),
183            );
184        }
185
186        cls = cls.add_method(build_to_dict_method(&all_fields, ir, &obj.properties));
187        cls = cls.add_method(build_from_dict_method(
188            &name,
189            &all_fields,
190            ir,
191            &obj.properties,
192        ));
193    }
194
195    file = file.add_type(cls.build().ok()?);
196    file.build().ok()
197}
198
199fn build_to_dict_method(
200    all_fields: &[(&String, &IrProperty)],
201    ir: &IrSpec,
202    properties: &indexmap::IndexMap<String, IrProperty>,
203) -> FunSpec {
204    let self_param = ParameterSpec::of("self", TypeName::primitive(""));
205    let return_type = TypeName::generic(
206        TypeName::primitive("dict"),
207        vec![TypeName::primitive("str"), TypeName::primitive("object")],
208    );
209
210    let mut body = CodeBlock::builder();
211    body.add_statement("result: dict[str, object] = {}", ());
212    for (json_name, prop) in all_fields {
213        let field_name = python_field_name(&prop.name);
214        let to_expr = render_to_dict_expr(&format!("self.{field_name}"), json_name, ir, properties);
215        if prop.required && !prop.nullable {
216            body.add_statement(&format!("result[\"{json_name}\"] = {to_expr}"), ());
217        } else {
218            body.add_statement(&format!("if self.{field_name} is not None:%>"), ());
219            body.add_statement(&format!("result[\"{json_name}\"] = {to_expr}%<"), ());
220        }
221    }
222    body.add_statement("return result", ());
223
224    FunSpec::builder("to_dict")
225        .add_param(self_param)
226        .returns(return_type)
227        .body(body.build().expect("to_dict body"))
228        .build()
229        .expect("to_dict method")
230}
231
232fn build_from_dict_method(
233    class_name: &str,
234    all_fields: &[(&String, &IrProperty)],
235    ir: &IrSpec,
236    properties: &indexmap::IndexMap<String, IrProperty>,
237) -> FunSpec {
238    let cls_param = ParameterSpec::of("cls", TypeName::primitive(""));
239    let data_param = ParameterSpec::of(
240        "data",
241        TypeName::generic(
242            TypeName::primitive("dict"),
243            vec![TypeName::primitive("str"), TypeName::primitive("object")],
244        ),
245    );
246
247    let mut body = CodeBlock::builder();
248    body.add_statement("return cls(%>", ());
249    for (json_name, prop) in all_fields {
250        let field_name = python_field_name(&prop.name);
251        let is_required = prop.required && !prop.nullable;
252        let expr = if is_required {
253            render_from_dict_expr(json_name, ir, properties)
254        } else {
255            render_from_dict_optional_expr(json_name, ir, properties)
256        };
257        if let Some(comment_start) = expr.find("  #") {
258            let (value_part, comment_part) = expr.split_at(comment_start);
259            body.add_statement(&format!("{field_name}={value_part},{comment_part}"), ());
260        } else {
261            body.add_statement(&format!("{field_name}={expr},"), ());
262        }
263    }
264    body.add("%<", ());
265    body.add_statement(")", ());
266
267    FunSpec::builder("from_dict")
268        .annotation(CodeBlock::of("@classmethod", ()).expect("classmethod"))
269        .add_param(cls_param)
270        .add_param(data_param)
271        .returns(TypeName::primitive(class_name))
272        .body(body.build().expect("from_dict body"))
273        .build()
274        .expect("from_dict method")
275}
276
277// ---------------------------------------------------------------------------
278// Enum -> class(str, Enum) or class(int, Enum)
279// ---------------------------------------------------------------------------
280
281fn emit_enum(schema: &IrSchema, en: &IrEnum) -> Option<FileSpec> {
282    if en.value_type == IrEnumValueType::Mixed {
283        return emit_type_alias_raw(schema, "object");
284    }
285
286    let name = schema.name.to_pascal_case();
287    let base = match en.value_type {
288        IrEnumValueType::String => TypeName::primitive("str"),
289        IrEnumValueType::Integer | IrEnumValueType::Number => TypeName::primitive("int"),
290        IrEnumValueType::Mixed => unreachable!(),
291    };
292
293    let mut ts = TypeSpec::builder(&name, TypeKind::Enum)
294        .extends(base)
295        .extends(TypeName::importable("enum", "Enum"));
296
297    if let Some(doc) = &schema.description {
298        ts = ts.doc(&format!("{}.", escape_docstring(doc)));
299    }
300
301    for v in &en.values {
302        let (member_name, value_code) = match en.value_type {
303            IrEnumValueType::String => {
304                let s = v.value.as_str()?;
305                (
306                    python_enum_member_name(s),
307                    format!("\"{}\"", escape_python_string(s)),
308                )
309            }
310            IrEnumValueType::Integer | IrEnumValueType::Number => {
311                let n = v
312                    .value
313                    .as_i64()
314                    .or_else(|| v.value.as_f64().map(|f| f as i64))?;
315                (format!("N{n}").replace('-', "NEG"), format!("{n}"))
316            }
317            IrEnumValueType::Mixed => unreachable!(),
318        };
319        ts = ts.add_variant(
320            EnumVariantSpec::builder(&member_name)
321                .value(CodeBlock::of(&value_code, ()).expect("enum value"))
322                .build()
323                .expect("enum variant"),
324        );
325    }
326
327    let file = FileSpec::builder_with("model.py", Python::new())
328        .header(future_annotations_header())
329        .add_type(ts.build().ok()?);
330    file.build().ok()
331}
332
333// ---------------------------------------------------------------------------
334// Alias -> type X = Y (PEP 695)
335// ---------------------------------------------------------------------------
336
337fn emit_alias(schema: &IrSchema, expr: &IrTypeExpr) -> Option<FileSpec> {
338    let name = schema.name.to_pascal_case();
339    let rhs_type = python_type_name(expr);
340
341    let type_alias = sigil_quote!(Python {
342        type $N(name.as_str()) = ($T(rhs_type));
343    })
344    .ok()?;
345
346    let mut file =
347        FileSpec::builder_with("model.py", Python::new()).header(future_annotations_header());
348    if needs_typing_literal(expr) {
349        file = file.add_import(ImportSpec::named("typing", "Literal"));
350    }
351    if let Some(doc) = &schema.description {
352        file = file.add_raw(&format!("# {}\n", escape_docstring(doc)));
353    }
354    file = file.add_code(type_alias);
355    file.build().ok()
356}
357
358fn emit_type_alias_raw(schema: &IrSchema, rhs: &str) -> Option<FileSpec> {
359    let name = schema.name.to_pascal_case();
360
361    let type_alias = sigil_quote!(Python {
362        type $N(name.as_str()) = $L(rhs);
363    })
364    .ok()?;
365
366    let mut file =
367        FileSpec::builder_with("model.py", Python::new()).header(future_annotations_header());
368    if let Some(doc) = &schema.description {
369        file = file.add_raw(&format!("# {}\n", escape_docstring(doc)));
370    }
371    file = file.add_code(type_alias);
372    file.build().ok()
373}
374
375// ---------------------------------------------------------------------------
376// Type alias builder with proper multi-line formatting
377// ---------------------------------------------------------------------------
378
379fn format_type_alias(name: &str, members: &[TypeName]) -> CodeBlock {
380    if members.is_empty() {
381        return sigil_quote!(Python {
382            type $N(name) = ($T(TypeName::importable("typing", "Any")));
383        })
384        .unwrap();
385    }
386    if members.len() == 1 {
387        return sigil_quote!(Python {
388            type $N(name) = ($T(members[0].clone()));
389        })
390        .unwrap();
391    }
392    sigil_quote!(Python {
393        type $N(name) = (
394            $L("    ")$for(member in members; separator = "\n    | ") { $T((*member).clone()) }
395        )
396    })
397    .unwrap()
398}
399
400// ---------------------------------------------------------------------------
401// Union -> type X = A | B | C
402// ---------------------------------------------------------------------------
403
404fn emit_union(schema: &IrSchema, u: &IrUnion) -> Option<FileSpec> {
405    let name = schema.name.to_pascal_case();
406
407    let mut members: Vec<TypeName> = u.members.iter().map(python_type_name).collect();
408    if u.nullable {
409        members.push(TypeName::primitive("None"));
410    }
411
412    let type_alias = format_type_alias(&name, &members);
413
414    let mut file =
415        FileSpec::builder_with("model.py", Python::new()).header(future_annotations_header());
416    if needs_typing_literal_in_exprs(&u.members) {
417        file = file.add_import(ImportSpec::named("typing", "Literal"));
418    }
419    if let Some(doc) = &schema.description {
420        file = file.add_raw(&format!("# {}\n", escape_docstring(doc)));
421    }
422    file = file.add_code(type_alias);
423    file.build().ok()
424}
425
426// ---------------------------------------------------------------------------
427// Intersection -> merged @dataclass
428// ---------------------------------------------------------------------------
429
430fn emit_intersection(schema: &IrSchema, inter: &IrIntersection, ir: &IrSpec) -> Option<FileSpec> {
431    let mut all_props: indexmap::IndexMap<String, IrProperty> = indexmap::IndexMap::new();
432    for member in &inter.members {
433        if let IrTypeExpr::Named(ref_name) = member
434            && let Some(s) = ir.schemas.get(ref_name.as_str())
435            && let IrSchemaKind::Object(obj) = &s.kind
436        {
437            for (k, v) in &obj.properties {
438                all_props.entry(k.clone()).or_insert_with(|| v.clone());
439            }
440        }
441    }
442
443    if all_props.is_empty() {
444        return emit_intersection_as_alias(schema, inter);
445    }
446
447    emit_intersection_as_dataclass(schema, &all_props, ir)
448}
449
450fn emit_intersection_as_alias(schema: &IrSchema, inter: &IrIntersection) -> Option<FileSpec> {
451    let name = schema.name.to_pascal_case();
452    let members: Vec<TypeName> = inter.members.iter().map(python_type_name).collect();
453
454    let type_alias = format_type_alias(&name, &members);
455
456    let mut file =
457        FileSpec::builder_with("model.py", Python::new()).header(future_annotations_header());
458    if needs_typing_literal_in_exprs(&inter.members) {
459        file = file.add_import(ImportSpec::named("typing", "Literal"));
460    }
461    file = file.add_code(type_alias);
462    file.build().ok()
463}
464
465fn emit_intersection_as_dataclass(
466    schema: &IrSchema,
467    all_props: &indexmap::IndexMap<String, IrProperty>,
468    ir: &IrSpec,
469) -> Option<FileSpec> {
470    let name = schema.name.to_pascal_case();
471
472    let mut file =
473        FileSpec::builder_with("model.py", Python::new()).header(future_annotations_header());
474
475    if needs_typing_literal_in_props(all_props) {
476        file = file.add_import(ImportSpec::named("typing", "Literal"));
477    }
478
479    let dataclass_tn = TypeName::importable("dataclasses", "dataclass");
480    let mut cls = TypeSpec::builder(&name, TypeKind::Class)
481        .annotate(AnnotationSpec::importable(dataclass_tn));
482
483    if let Some(doc) = &schema.description {
484        cls = cls.doc(&format!("{}.", escape_docstring(doc)));
485    }
486
487    let mut required: Vec<(&String, &IrProperty)> = Vec::new();
488    let mut optional: Vec<(&String, &IrProperty)> = Vec::new();
489    for (json_name, prop) in all_props {
490        if prop.required && !prop.nullable {
491            required.push((json_name, prop));
492        } else {
493            optional.push((json_name, prop));
494        }
495    }
496
497    if required.is_empty() && optional.is_empty() {
498        cls = cls.extra_member(CodeBlock::of("pass", ()).expect("pass"));
499    } else {
500        for (_json_name, prop) in &required {
501            let field_name = python_field_name(&prop.name);
502            let type_name = python_type_name_for_schema(&prop.type_expr, &schema.name);
503            cls = cls.add_field(
504                FieldSpec::builder(&field_name, type_name)
505                    .build()
506                    .expect("required field"),
507            );
508        }
509        for (_json_name, prop) in &optional {
510            let field_name = python_field_name(&prop.name);
511            let type_name = python_type_name_for_schema(&prop.type_expr, &schema.name);
512            cls = cls.add_field(
513                FieldSpec::builder(&field_name, TypeName::optional(type_name))
514                    .initializer(CodeBlock::of("None", ()).expect("None init"))
515                    .build()
516                    .expect("optional field"),
517            );
518        }
519
520        let all_fields: Vec<(&String, &IrProperty)> =
521            required.iter().chain(optional.iter()).copied().collect();
522        cls = cls.add_method(build_to_dict_method(&all_fields, ir, all_props));
523        cls = cls.add_method(build_from_dict_method(&name, &all_fields, ir, all_props));
524    }
525
526    file = file.add_type(cls.build().ok()?);
527    file.build().ok()
528}
529
530// ---------------------------------------------------------------------------
531// TaggedUnion -> type X = A | B | C
532// ---------------------------------------------------------------------------
533
534fn emit_tagged_union(schema: &IrSchema, tu: &IrTaggedUnion, ir: &IrSpec) -> Option<FileSpec> {
535    let name = schema.name.to_pascal_case();
536    let snake_name = schema.name.to_snake_case();
537
538    let members: Vec<TypeName> = tu
539        .variants
540        .iter()
541        .map(|v| python_type_name(&v.content_type))
542        .collect();
543
544    let type_alias = format_type_alias(&name, &members);
545
546    let hint = match &tu.tagging {
547        TaggingStyle::Internal => {
548            format!("Discriminator: {} (internal).", tu.discriminator_field)
549        }
550        TaggingStyle::Adjacent { content_field } => format!(
551            "Discriminator: {} / content: {} (adjacent).",
552            tu.discriminator_field, content_field
553        ),
554        TaggingStyle::External => "Discriminator: variant key (external).".to_string(),
555    };
556
557    let doc = match &schema.description {
558        Some(desc) => format!("{desc}\n\n{hint}"),
559        None => hint,
560    };
561
562    let mut file =
563        FileSpec::builder_with("model.py", Python::new()).header(future_annotations_header());
564    let exprs: Vec<&IrTypeExpr> = tu.variants.iter().map(|v| &v.content_type).collect();
565    if exprs.iter().any(|e| needs_typing_literal(e)) {
566        file = file.add_import(ImportSpec::named("typing", "Literal"));
567    }
568    let mut doc_block = String::new();
569    for line in doc.lines() {
570        doc_block.push_str(&format!("# {line}\n"));
571    }
572    file = file.add_raw(&doc_block);
573    file = file.add_code(type_alias);
574
575    if !tu.variants.is_empty() {
576        let helpers = build_tagged_union_helpers(&name, &snake_name, tu, ir);
577        file = file.add_code(helpers);
578    }
579
580    file.build().ok()
581}
582
583fn build_tagged_union_helpers(
584    pascal_name: &str,
585    snake_name: &str,
586    tu: &IrTaggedUnion,
587    ir: &IrSpec,
588) -> CodeBlock {
589    let tag_field = &tu.discriminator_field;
590
591    // Only generate helpers for variants that resolve to Object schemas
592    let resolved_variants: Vec<(&IrTaggedVariant, String)> = tu
593        .variants
594        .iter()
595        .filter_map(|v| {
596            if let IrTypeExpr::Named(ref_name) = &v.content_type
597                && is_object_schema(ref_name, ir)
598            {
599                return Some((v, ref_name.to_pascal_case()));
600            }
601            None
602        })
603        .collect();
604
605    let mut cb = CodeBlock::builder();
606
607    if resolved_variants.is_empty() {
608        return cb.build_unwrap();
609    }
610
611    // from_dict
612    cb.add_line();
613    cb.begin_control_flow(
614        &format!("def {snake_name}_from_dict(data: dict[str, object]) -> {pascal_name}"),
615        (),
616    );
617    match &tu.tagging {
618        TaggingStyle::Internal => {
619            cb.add_statement(&format!("_tag = data[\"{tag_field}\"]"), ());
620            for (i, (variant, py_class)) in resolved_variants.iter().enumerate() {
621                let cond = format!("_tag == \"{}\"", variant.discriminator_value);
622                emit_elif(&mut cb, i == 0, false, &cond);
623                cb.add_statement(&format!("return {py_class}.from_dict(data)"), ());
624            }
625            cb.end_control_flow_no_newline();
626        }
627        TaggingStyle::Adjacent { content_field } => {
628            cb.add_statement(&format!("_tag = data[\"{tag_field}\"]"), ());
629            cb.add_statement(
630                &format!("_content = data[\"{content_field}\"]  # type: ignore[assignment]"),
631                (),
632            );
633            for (i, (variant, py_class)) in resolved_variants.iter().enumerate() {
634                let cond = format!("_tag == \"{}\"", variant.discriminator_value);
635                emit_elif(&mut cb, i == 0, false, &cond);
636                cb.add_statement(
637                    &format!("return {py_class}.from_dict(_content)  # type: ignore[arg-type]"),
638                    (),
639                );
640            }
641            cb.end_control_flow_no_newline();
642        }
643        TaggingStyle::External => {
644            for (i, (variant, py_class)) in resolved_variants.iter().enumerate() {
645                let cond = format!("\"{}\" in data", variant.discriminator_value);
646                emit_elif(&mut cb, i == 0, false, &cond);
647                cb.add_statement(
648                    &format!(
649                        "return {py_class}.from_dict(data[\"{}\"])  # type: ignore[arg-type]",
650                        variant.discriminator_value
651                    ),
652                    (),
653                );
654            }
655            cb.end_control_flow_no_newline();
656        }
657    }
658    cb.add_statement(
659        "raise ValueError(%V)",
660        VerbatimStrArg(format!(
661            "Unknown discriminator value for {pascal_name}: {{data}}"
662        )),
663    );
664    cb.end_control_flow();
665
666    // to_dict
667    cb.add_line();
668    cb.begin_control_flow(
669        &format!("def {snake_name}_to_dict(obj: {pascal_name}) -> dict[str, object]"),
670        (),
671    );
672    match &tu.tagging {
673        TaggingStyle::Internal => {
674            for (i, (variant, py_class)) in resolved_variants.iter().enumerate() {
675                let cond = format!("isinstance(obj, {py_class})");
676                emit_elif(&mut cb, i == 0, false, &cond);
677                cb.add_statement("result = obj.to_dict()", ());
678                cb.add_statement(
679                    &format!(
680                        "result[\"{tag_field}\"] = \"{}\"",
681                        variant.discriminator_value
682                    ),
683                    (),
684                );
685                cb.add_statement("return result", ());
686            }
687            cb.end_control_flow_no_newline();
688        }
689        TaggingStyle::Adjacent { content_field } => {
690            for (i, (variant, py_class)) in resolved_variants.iter().enumerate() {
691                let cond = format!("isinstance(obj, {py_class})");
692                emit_elif(&mut cb, i == 0, false, &cond);
693                cb.add_statement(
694                    &format!(
695                        "return {{\"{tag_field}\": \"{}\", \"{content_field}\": obj.to_dict()}}",
696                        variant.discriminator_value
697                    ),
698                    (),
699                );
700            }
701            cb.end_control_flow_no_newline();
702        }
703        TaggingStyle::External => {
704            for (i, (variant, py_class)) in resolved_variants.iter().enumerate() {
705                let cond = format!("isinstance(obj, {py_class})");
706                emit_elif(&mut cb, i == 0, false, &cond);
707                cb.add_statement(
708                    &format!(
709                        "return {{\"{}\": obj.to_dict()}}",
710                        variant.discriminator_value
711                    ),
712                    (),
713                );
714            }
715            cb.end_control_flow_no_newline();
716        }
717    }
718    cb.add_statement(
719        "raise ValueError(%V)",
720        VerbatimStrArg(format!("Unknown variant for {pascal_name}: {{type(obj)}}")),
721    );
722    cb.end_control_flow();
723
724    cb.build_unwrap()
725}
726
727fn emit_elif(
728    cb: &mut sigil_stitch::code_block::CodeBlockBuilder,
729    is_first: bool,
730    is_last: bool,
731    cond: &str,
732) {
733    if !is_first {
734        cb.end_control_flow_no_newline();
735    }
736    if is_last && !is_first {
737        cb.begin_control_flow("else", ());
738    } else {
739        let kw = if is_first { "if" } else { "elif" };
740        cb.begin_control_flow(&format!("{kw} {cond}"), ());
741    }
742}
743
744// ---------------------------------------------------------------------------
745// Type mapping
746// ---------------------------------------------------------------------------
747
748/// Like `python_type_name` but Named types import from `..models.{snake}` (for API files).
749pub fn api_type_name(expr: &IrTypeExpr) -> TypeName {
750    match expr {
751        IrTypeExpr::Named(name) => {
752            let py_name = name.to_pascal_case();
753            let module = format!("..models.{}", name.to_snake_case());
754            TypeName::importable(&module, &py_name)
755        }
756        IrTypeExpr::Array(inner) => {
757            TypeName::generic(TypeName::primitive("list"), vec![api_type_name(inner)])
758        }
759        IrTypeExpr::Map(inner) => TypeName::generic(
760            TypeName::primitive("dict"),
761            vec![TypeName::primitive("str"), api_type_name(inner)],
762        ),
763        IrTypeExpr::Union(members) => {
764            if members.is_empty() {
765                TypeName::importable("typing", "Any")
766            } else {
767                TypeName::union(members.iter().map(api_type_name).collect())
768            }
769        }
770        IrTypeExpr::Nullable(inner) => TypeName::optional(api_type_name(inner)),
771        _ => python_type_name(expr),
772    }
773}
774
775/// Map an IR type expression to a Python type string (for serialization helpers).
776pub fn python_type_str(expr: &IrTypeExpr) -> String {
777    match expr {
778        IrTypeExpr::Named(name) => name.to_pascal_case(),
779        IrTypeExpr::Primitive(p) => python_primitive(p).to_string(),
780        IrTypeExpr::StringLiteral(s) => {
781            format!("Literal[\"{}\"]", escape_python_string(s))
782        }
783        IrTypeExpr::StringEnum(values) => {
784            let members: Vec<String> = values
785                .iter()
786                .map(|v| format!("\"{}\"", escape_python_string(v)))
787                .collect();
788            format!("Literal[{}]", members.join(", "))
789        }
790        IrTypeExpr::Array(inner) => {
791            let inner_ty = python_type_str(inner);
792            format!("list[{inner_ty}]")
793        }
794        IrTypeExpr::Map(inner) => {
795            let inner_ty = python_type_str(inner);
796            format!("dict[str, {inner_ty}]")
797        }
798        IrTypeExpr::Union(members) => {
799            let parts: Vec<String> = members.iter().map(python_type_str).collect();
800            if parts.is_empty() {
801                "Any".to_string()
802            } else {
803                parts.join(" | ")
804            }
805        }
806        IrTypeExpr::Nullable(inner) => {
807            let inner_ty = python_type_str(inner);
808            format!("{inner_ty} | None")
809        }
810        IrTypeExpr::Any => "Any".to_string(),
811    }
812}
813
814fn python_primitive(p: &IrPrimitive) -> &'static str {
815    match p {
816        IrPrimitive::String | IrPrimitive::StringWithFormat(_) => "str",
817        IrPrimitive::Integer | IrPrimitive::IntegerWithFormat(_) => "int",
818        IrPrimitive::Number | IrPrimitive::NumberWithFormat(_) => "float",
819        IrPrimitive::Boolean => "bool",
820        IrPrimitive::Binary => "bytes",
821        IrPrimitive::Date => "datetime.date",
822        IrPrimitive::DateTime => "datetime.datetime",
823        IrPrimitive::Uuid => "uuid.UUID",
824    }
825}
826
827fn collect_request_input_imports(
828    expr: &IrTypeExpr,
829    imports: &mut std::collections::BTreeSet<String>,
830) {
831    match expr {
832        IrTypeExpr::Named(name) => {
833            let py_name = name.to_pascal_case();
834            let module = name.to_snake_case();
835            imports.insert(format!("from .{module} import {py_name}"));
836        }
837        IrTypeExpr::Primitive(IrPrimitive::Date | IrPrimitive::DateTime) => {
838            imports.insert("import datetime".to_string());
839        }
840        IrTypeExpr::Primitive(IrPrimitive::Uuid) => {
841            imports.insert("import uuid".to_string());
842        }
843        IrTypeExpr::StringLiteral(_) | IrTypeExpr::StringEnum(_) => {
844            imports.insert("from typing import Literal".to_string());
845        }
846        IrTypeExpr::Union(members) => {
847            if members.is_empty() {
848                imports.insert("from typing import Any".to_string());
849            }
850            for member in members {
851                collect_request_input_imports(member, imports);
852            }
853        }
854        IrTypeExpr::Any => {
855            imports.insert("from typing import Any".to_string());
856        }
857        IrTypeExpr::Array(inner) | IrTypeExpr::Map(inner) | IrTypeExpr::Nullable(inner) => {
858            collect_request_input_imports(inner, imports);
859        }
860        _ => {}
861    }
862}
863
864fn needs_typing_literal(expr: &IrTypeExpr) -> bool {
865    match expr {
866        IrTypeExpr::StringLiteral(_) | IrTypeExpr::StringEnum(_) => true,
867        IrTypeExpr::Array(inner) | IrTypeExpr::Map(inner) | IrTypeExpr::Nullable(inner) => {
868            needs_typing_literal(inner)
869        }
870        IrTypeExpr::Union(members) => members.iter().any(needs_typing_literal),
871        _ => false,
872    }
873}
874
875fn needs_typing_literal_in_props(props: &indexmap::IndexMap<String, IrProperty>) -> bool {
876    props.values().any(|p| needs_typing_literal(&p.type_expr))
877}
878
879fn needs_typing_literal_in_exprs(exprs: &[IrTypeExpr]) -> bool {
880    exprs.iter().any(needs_typing_literal)
881}
882
883// ---------------------------------------------------------------------------
884// Serialization helpers
885// ---------------------------------------------------------------------------
886
887fn render_to_dict_expr(
888    value_expr: &str,
889    json_name: &str,
890    ir: &IrSpec,
891    properties: &indexmap::IndexMap<String, IrProperty>,
892) -> String {
893    let prop = properties.get(json_name);
894    let type_expr = prop.map(|p| &p.type_expr);
895    match type_expr {
896        Some(IrTypeExpr::Named(ref_name)) => {
897            if is_object_schema(ref_name, ir) {
898                format!("{value_expr}.to_dict()")
899            } else if is_tagged_union_schema(ref_name, ir) {
900                let snake = ref_name.to_snake_case();
901                format!("{snake}_to_dict({value_expr})")
902            } else {
903                value_expr.to_string()
904            }
905        }
906        Some(IrTypeExpr::Array(inner)) => {
907            if let IrTypeExpr::Named(ref_name) = inner.as_ref()
908                && is_object_schema(ref_name, ir)
909            {
910                return format!("[item.to_dict() for item in {value_expr}]");
911            }
912            if let IrTypeExpr::Named(ref_name) = inner.as_ref()
913                && is_tagged_union_schema(ref_name, ir)
914            {
915                let snake = ref_name.to_snake_case();
916                return format!("[{snake}_to_dict(item) for item in {value_expr}]");
917            }
918            value_expr.to_string()
919        }
920        Some(IrTypeExpr::Nullable(inner)) => {
921            if let IrTypeExpr::Named(ref_name) = inner.as_ref()
922                && is_object_schema(ref_name, ir)
923            {
924                return format!("{value_expr}.to_dict() if {value_expr} is not None else None");
925            }
926            if let IrTypeExpr::Named(ref_name) = inner.as_ref()
927                && is_tagged_union_schema(ref_name, ir)
928            {
929                let snake = ref_name.to_snake_case();
930                return format!(
931                    "{snake}_to_dict({value_expr}) if {value_expr} is not None else None"
932                );
933            }
934            value_expr.to_string()
935        }
936        Some(IrTypeExpr::Map(inner)) => {
937            if let IrTypeExpr::Named(ref_name) = inner.as_ref()
938                && is_object_schema(ref_name, ir)
939            {
940                return format!("{{k: v.to_dict() for k, v in {value_expr}.items()}}");
941            }
942            if let IrTypeExpr::Named(ref_name) = inner.as_ref()
943                && is_tagged_union_schema(ref_name, ir)
944            {
945                let snake = ref_name.to_snake_case();
946                return format!("{{k: {snake}_to_dict(v) for k, v in {value_expr}.items()}}");
947            }
948            value_expr.to_string()
949        }
950        _ => value_expr.to_string(),
951    }
952}
953
954fn render_from_dict_expr(
955    json_name: &str,
956    ir: &IrSpec,
957    properties: &indexmap::IndexMap<String, IrProperty>,
958) -> String {
959    let prop = properties.get(json_name);
960    let type_expr = prop.map(|p| &p.type_expr);
961    let accessor = format!("data[\"{json_name}\"]");
962    match type_expr {
963        Some(IrTypeExpr::Named(ref_name)) => {
964            if is_object_schema(ref_name, ir) {
965                let py_name = ref_name.to_pascal_case();
966                format!("{py_name}.from_dict({accessor})  # type: ignore[arg-type]")
967            } else if is_tagged_union_schema(ref_name, ir) {
968                let snake = ref_name.to_snake_case();
969                format!("{snake}_from_dict({accessor})  # type: ignore[arg-type]")
970            } else {
971                format!("{accessor}  # type: ignore[assignment]")
972            }
973        }
974        Some(IrTypeExpr::Array(inner)) => {
975            if let IrTypeExpr::Named(ref_name) = inner.as_ref()
976                && is_object_schema(ref_name, ir)
977            {
978                let py_name = ref_name.to_pascal_case();
979                return format!(
980                    "[{py_name}.from_dict(item) for item in {accessor}]  # type: ignore[union-attr]"
981                );
982            }
983            if let IrTypeExpr::Named(ref_name) = inner.as_ref()
984                && is_tagged_union_schema(ref_name, ir)
985            {
986                let snake = ref_name.to_snake_case();
987                return format!(
988                    "[{snake}_from_dict(item) for item in {accessor}]  # type: ignore[union-attr]"
989                );
990            }
991            format!("{accessor}  # type: ignore[assignment]")
992        }
993        Some(IrTypeExpr::Map(inner)) => {
994            if let IrTypeExpr::Named(ref_name) = inner.as_ref()
995                && is_object_schema(ref_name, ir)
996            {
997                let py_name = ref_name.to_pascal_case();
998                return format!(
999                    "{{k: {py_name}.from_dict(v) for k, v in {accessor}.items()}}  # type: ignore[union-attr]"
1000                );
1001            }
1002            if let IrTypeExpr::Named(ref_name) = inner.as_ref()
1003                && is_tagged_union_schema(ref_name, ir)
1004            {
1005                let snake = ref_name.to_snake_case();
1006                return format!(
1007                    "{{k: {snake}_from_dict(v) for k, v in {accessor}.items()}}  # type: ignore[union-attr]"
1008                );
1009            }
1010            format!("{accessor}  # type: ignore[assignment]")
1011        }
1012        _ => format!("{accessor}  # type: ignore[assignment]"),
1013    }
1014}
1015
1016fn render_from_dict_optional_expr(
1017    json_name: &str,
1018    ir: &IrSpec,
1019    properties: &indexmap::IndexMap<String, IrProperty>,
1020) -> String {
1021    let prop = properties.get(json_name);
1022    let type_expr = prop.map(|p| &p.type_expr);
1023    let raw_type = type_expr.map(|t| match t {
1024        IrTypeExpr::Nullable(inner) => inner.as_ref(),
1025        _ => t,
1026    });
1027    let accessor = format!("data.get(\"{json_name}\")");
1028    match raw_type {
1029        Some(IrTypeExpr::Named(ref_name)) => {
1030            if is_object_schema(ref_name, ir) {
1031                let py_name = ref_name.to_pascal_case();
1032                format!(
1033                    "{py_name}.from_dict({accessor}) if {accessor} is not None else None  # type: ignore[arg-type]"
1034                )
1035            } else if is_tagged_union_schema(ref_name, ir) {
1036                let snake = ref_name.to_snake_case();
1037                format!(
1038                    "{snake}_from_dict({accessor}) if {accessor} is not None else None  # type: ignore[arg-type]"
1039                )
1040            } else {
1041                format!("{accessor}  # type: ignore[assignment]")
1042            }
1043        }
1044        Some(IrTypeExpr::Array(inner)) => {
1045            if let IrTypeExpr::Named(ref_name) = inner.as_ref()
1046                && is_object_schema(ref_name, ir)
1047            {
1048                let py_name = ref_name.to_pascal_case();
1049                return format!(
1050                    "[{py_name}.from_dict(item) for item in {accessor}] if {accessor} is not None else None  # type: ignore[union-attr]"
1051                );
1052            }
1053            if let IrTypeExpr::Named(ref_name) = inner.as_ref()
1054                && is_tagged_union_schema(ref_name, ir)
1055            {
1056                let snake = ref_name.to_snake_case();
1057                return format!(
1058                    "[{snake}_from_dict(item) for item in {accessor}] if {accessor} is not None else None  # type: ignore[union-attr]"
1059                );
1060            }
1061            format!("{accessor}  # type: ignore[assignment]")
1062        }
1063        Some(IrTypeExpr::Map(inner)) => {
1064            if let IrTypeExpr::Named(ref_name) = inner.as_ref()
1065                && is_object_schema(ref_name, ir)
1066            {
1067                let py_name = ref_name.to_pascal_case();
1068                return format!(
1069                    "{{k: {py_name}.from_dict(v) for k, v in {accessor}.items()}} if {accessor} is not None else None  # type: ignore[union-attr]"
1070                );
1071            }
1072            if let IrTypeExpr::Named(ref_name) = inner.as_ref()
1073                && is_tagged_union_schema(ref_name, ir)
1074            {
1075                let snake = ref_name.to_snake_case();
1076                return format!(
1077                    "{{k: {snake}_from_dict(v) for k, v in {accessor}.items()}} if {accessor} is not None else None  # type: ignore[union-attr]"
1078                );
1079            }
1080            format!("{accessor}  # type: ignore[assignment]")
1081        }
1082        _ => format!("{accessor}  # type: ignore[assignment]"),
1083    }
1084}
1085
1086pub fn is_object_schema(name: &str, ir: &IrSpec) -> bool {
1087    ir.schemas.get(name).is_some_and(|s| match &s.kind {
1088        IrSchemaKind::Object(_) => true,
1089        IrSchemaKind::Intersection(inter) => inter.members.iter().any(|m| {
1090            if let IrTypeExpr::Named(ref_name) = m {
1091                ir.schemas
1092                    .get(ref_name.as_str())
1093                    .is_some_and(|ms| matches!(ms.kind, IrSchemaKind::Object(_)))
1094            } else {
1095                false
1096            }
1097        }),
1098        _ => false,
1099    })
1100}
1101
1102pub fn is_tagged_union_schema(name: &str, ir: &IrSpec) -> bool {
1103    ir.schemas
1104        .get(name)
1105        .is_some_and(|s| matches!(s.kind, IrSchemaKind::TaggedUnion(_)))
1106}
1107
1108/// Collect all tagged-union schema names referenced (directly or nested) in a type expression.
1109fn collect_tagged_union_refs(expr: &IrTypeExpr, ir: &IrSpec) -> Vec<String> {
1110    let mut refs = Vec::new();
1111    match expr {
1112        IrTypeExpr::Named(name) => {
1113            if is_tagged_union_schema(name, ir) {
1114                refs.push(name.clone());
1115            }
1116        }
1117        IrTypeExpr::Array(inner) | IrTypeExpr::Nullable(inner) => {
1118            refs.extend(collect_tagged_union_refs(inner, ir));
1119        }
1120        IrTypeExpr::Map(inner) => {
1121            refs.extend(collect_tagged_union_refs(inner, ir));
1122        }
1123        _ => {}
1124    }
1125    refs
1126}
1127
1128// ---------------------------------------------------------------------------
1129// Helpers
1130// ---------------------------------------------------------------------------
1131
1132pub fn python_field_name(name: &str) -> String {
1133    let snake = name.to_snake_case();
1134    if snake.is_empty() {
1135        return "field_".to_string();
1136    }
1137    match snake.as_str() {
1138        "and" | "as" | "assert" | "async" | "await" | "break" | "class" | "continue" | "def"
1139        | "del" | "elif" | "else" | "except" | "finally" | "for" | "from" | "global" | "if"
1140        | "import" | "in" | "is" | "lambda" | "nonlocal" | "not" | "or" | "pass" | "raise"
1141        | "return" | "try" | "while" | "with" | "yield" | "type" => {
1142            format!("{snake}_")
1143        }
1144        _ => snake,
1145    }
1146}
1147
1148fn python_enum_member_name(value: &str) -> String {
1149    let upper = value
1150        .to_uppercase()
1151        .replace(|c: char| !c.is_alphanumeric(), "_");
1152    if upper.is_empty() {
1153        return "EMPTY".to_string();
1154    }
1155    if upper.starts_with(|c: char| c.is_ascii_digit()) {
1156        return format!("N{upper}");
1157    }
1158    upper
1159}
1160
1161fn escape_python_string(s: &str) -> String {
1162    s.replace('\\', "\\\\").replace('"', "\\\"")
1163}
1164
1165fn escape_docstring(s: &str) -> String {
1166    s.replace("\"\"\"", "\\\"\\\"\\\"")
1167        .lines()
1168        .next()
1169        .unwrap_or("")
1170        .to_string()
1171}
openapi_nexus/generators/python/httpx/emit_models.rs

openapi_nexus/generators/python/httpx/
emit_models.rs
