pushkin-compiler 0.2.1

//! Pydantic v2 emission: `ConfigDict(extra='forbid', strict=True)` (spec §5.2).

use std::fmt::Write as _;

use crate::schema::{ArrayElement, ContractSchema, PropertyKind};
use crate::targets::{header, python_number_literal, quoted_string, type_name};

pub fn emit(schema: &ContractSchema, epoch: u32) -> String {
    let mut fields = String::new();
    let mut needs_literal = false;
    let mut needs_optional = false;

    for property in &schema.properties {
        let field = python_field(&property.kind);
        needs_literal |= field.needs_literal;

        let mut type_annotation = field.type_annotation;
        let optional = !property.required && !property.kind.has_default();
        if optional {
            needs_optional = true;
            type_annotation = format!("Optional[{type_annotation}]");
        }

        let mut constraints = field.constraints;
        if let Some(literal) = field.default_literal {
            constraints.push(format!("default={literal}"));
        }

        let field_call = if constraints.is_empty() && optional {
            " = None".to_owned()
        } else if constraints.is_empty() {
            String::new()
        } else {
            if optional {
                constraints.insert(0, "default=None".to_owned());
            }
            format!(" = Field({})", constraints.join(", "))
        };
        let _ = writeln!(
            fields,
            "    {}: {type_annotation}{field_call}",
            property.name
        );
    }

    let mut typing_names = Vec::new();
    if needs_literal {
        typing_names.push("Literal");
    }
    if needs_optional {
        typing_names.push("Optional");
    }
    let typing_import = if typing_names.is_empty() {
        String::new()
    } else {
        format!("from typing import {}\n\n", typing_names.join(", "))
    };
    let name = type_name(schema);
    format!(
        "{header}{typing_import}from pydantic import BaseModel, ConfigDict, EmailStr, Field\n\n\n\
         class {name}(BaseModel):\n    model_config = ConfigDict(extra='forbid', strict=True)\n\n{fields}",
        header = header(schema, epoch, "#"),
    )
}

struct PythonField {
    type_annotation: String,
    constraints: Vec<String>,
    default_literal: Option<String>,
    needs_literal: bool,
}

/// Annotation, `Field()` constraints, and the already-spelled default
/// literal per kind. Python spells booleans `True`/`False`.
fn python_field(kind: &PropertyKind) -> PythonField {
    match kind {
        PropertyKind::String {
            min_length,
            max_length,
            format,
            enum_values,
            default,
        } => {
            let type_annotation = match enum_values {
                Some(values) => {
                    let list = values
                        .iter()
                        .map(|value| quoted_string(value))
                        .collect::<Vec<_>>()
                        .join(", ");
                    format!("Literal[{list}]")
                }
                None if format.as_deref() == Some("email") => "EmailStr".to_owned(),
                None => "str".to_owned(),
            };
            let mut constraints = Vec::new();
            if enum_values.is_none() {
                if let Some(min) = min_length {
                    constraints.push(format!("min_length={min}"));
                }
                if let Some(max) = max_length {
                    constraints.push(format!("max_length={max}"));
                }
            }
            PythonField {
                needs_literal: enum_values.is_some(),
                type_annotation,
                constraints,
                default_literal: default.as_deref().map(quoted_string),
            }
        }
        PropertyKind::Integer { default } => scalar_field("int", default.map(|v| v.to_string())),
        PropertyKind::Number { default } => {
            scalar_field("float", default.map(python_number_literal))
        }
        PropertyKind::Boolean { default } => scalar_field(
            "bool",
            default.map(|value| if value { "True" } else { "False" }.to_owned()),
        ),
        // D3: PEP 585 builtin generics — no `typing` import, so the
        // emitter's typing_names machinery stays exactly as it is. The
        // effective floor is already Python 3.9+ via the existing
        // `Optional` usage, which is the same floor PEP 585 requires.
        PropertyKind::Array { element } => scalar_field(
            &format!(
                "list[{}]",
                match element {
                    ArrayElement::String => "str",
                    ArrayElement::Integer => "int",
                    ArrayElement::Number => "float",
                    ArrayElement::Boolean => "bool",
                }
            ),
            None,
        ),
    }
}

fn scalar_field(annotation: &str, default_literal: Option<String>) -> PythonField {
    PythonField {
        type_annotation: annotation.to_owned(),
        constraints: Vec::new(),
        default_literal,
        needs_literal: false,
    }
}