use std::fmt::Write as _;
use crate::schema::{ContractSchema, PropertyKind};
use crate::targets::{header, 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 PropertyKind::String {
min_length,
max_length,
format,
enum_values,
default,
} = &property.kind;
let mut type_annotation = match enum_values {
Some(values) => {
needs_literal = true;
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 optional = !property.required && default.is_none();
if optional {
needs_optional = true;
type_annotation = format!("Optional[{type_annotation}]");
}
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}"));
}
}
if let Some(value) = default {
constraints.push(format!("default={}", quoted_string(value)));
}
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, "#"),
)
}