mold_cli/generators/zod/
mod.rs

1mod types;
2
3use crate::generators::{Generator, GeneratorConfig};
4use crate::types::{NestedType, ObjectType, Schema, SchemaType};
5use anyhow::Result;
6use std::collections::HashMap;
7
8use types::{format_field_name, generate_type};
9
10pub struct ZodGenerator;
11
12impl ZodGenerator {
13    pub fn new() -> Self {
14        Self
15    }
16
17    fn generate_schema(
18        &self,
19        name: &str,
20        obj: &ObjectType,
21        indent: &str,
22        type_refs: &HashMap<String, String>,
23        strict_mode: bool,
24    ) -> String {
25        let schema_name = format!("{}Schema", name);
26        let inner_indent = format!("{}  ", indent);
27        let mut lines = vec![format!("const {} = z.object({{", schema_name)];
28
29        for field in &obj.fields {
30            let field_name = format_field_name(&field.name);
31            let mut field_type = generate_type(&field.field_type, &inner_indent, type_refs);
32            if field.optional {
33                field_type = format!("{}.optional()", field_type);
34            }
35            lines.push(format!("{}{}: {},", inner_indent, field_name, field_type));
36        }
37
38        let strict_suffix = if strict_mode { ".strict()" } else { "" };
39        lines.push(format!("}}){};\n", strict_suffix));
40        lines.join("\n")
41    }
42
43    fn build_type_refs(&self, nested_types: &[NestedType]) -> HashMap<String, String> {
44        let mut refs = HashMap::new();
45        for nt in nested_types {
46            let key = format!("{:?}", nt.object);
47            refs.insert(key, nt.name.clone());
48        }
49        refs
50    }
51}
52
53impl Default for ZodGenerator {
54    fn default() -> Self {
55        Self::new()
56    }
57}
58
59impl Generator for ZodGenerator {
60    fn generate(&self, schema: &Schema, config: &GeneratorConfig) -> Result<String> {
61        let mut output = vec![
62            "// Generated by mold".to_string(),
63            "import { z } from \"zod\";".to_string(),
64            String::new(),
65        ];
66
67        let type_refs = if config.flat_mode {
68            HashMap::new()
69        } else {
70            self.build_type_refs(&schema.nested_types)
71        };
72
73        let mut all_type_names: Vec<String> = Vec::new();
74
75        if !config.flat_mode && !schema.nested_types.is_empty() {
76            for nt in schema.nested_types.iter().rev() {
77                output.push(self.generate_schema(
78                    &nt.name,
79                    &nt.object,
80                    &config.indent,
81                    &type_refs,
82                    config.zod_strict_objects,
83                ));
84                all_type_names.push(nt.name.clone());
85            }
86        }
87
88        if let SchemaType::Object(obj) = &schema.root_type {
89            output.push(self.generate_schema(
90                &schema.name,
91                obj,
92                &config.indent,
93                &type_refs,
94                config.zod_strict_objects,
95            ));
96            all_type_names.push(schema.name.clone());
97        }
98
99        for name in &all_type_names {
100            output.push(format!(
101                "type {} = z.infer<typeof {}Schema>;",
102                name, name
103            ));
104        }
105
106        output.push(String::new());
107
108        let schema_exports: Vec<String> =
109            all_type_names.iter().map(|n| format!("{}Schema", n)).collect();
110        output.push(format!("export {{ {} }};", schema_exports.join(", ")));
111        output.push(format!("export type {{ {} }};", all_type_names.join(", ")));
112
113        Ok(output.join("\n"))
114    }
115
116    fn file_extension(&self) -> &'static str {
117        "zod.ts"
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124    use crate::types::Field;
125
126    #[test]
127    fn test_generate_simple_schema() {
128        let gen = ZodGenerator::new();
129        let obj = ObjectType::new(vec![
130            Field::new("id", SchemaType::Integer),
131            Field::new("name", SchemaType::String),
132            Field::new("active", SchemaType::Boolean),
133        ]);
134        let schema = Schema::new("User", SchemaType::Object(obj));
135        let config = GeneratorConfig::default();
136
137        let output = gen.generate(&schema, &config).unwrap();
138
139        assert!(output.contains("import { z } from \"zod\""));
140        assert!(output.contains("const UserSchema = z.object({"));
141        assert!(output.contains("id: z.number().int()"));
142        assert!(output.contains("name: z.string()"));
143        assert!(output.contains("active: z.boolean()"));
144        assert!(output.contains("type User = z.infer<typeof UserSchema>"));
145        assert!(output.contains("export { UserSchema }"));
146    }
147
148    #[test]
149    fn test_generate_array_schema() {
150        let gen = ZodGenerator::new();
151        let obj = ObjectType::new(vec![Field::new(
152            "tags",
153            SchemaType::Array(Box::new(SchemaType::String)),
154        )]);
155        let schema = Schema::new("Test", SchemaType::Object(obj));
156        let config = GeneratorConfig::default();
157
158        let output = gen.generate(&schema, &config).unwrap();
159
160        assert!(output.contains("tags: z.array(z.string())"));
161    }
162
163    #[test]
164    fn test_generate_union_schema() {
165        let gen = ZodGenerator::new();
166        let obj = ObjectType::new(vec![Field::new(
167            "value",
168            SchemaType::Union(vec![SchemaType::String, SchemaType::Integer]),
169        )]);
170        let schema = Schema::new("Test", SchemaType::Object(obj));
171        let config = GeneratorConfig::default();
172
173        let output = gen.generate(&schema, &config).unwrap();
174
175        assert!(output.contains("z.union([z.string(), z.number().int()])"));
176    }
177
178    #[test]
179    fn test_empty_object() {
180        let gen = ZodGenerator::new();
181        let obj = ObjectType::new(vec![Field::new(
182            "metadata",
183            SchemaType::Object(ObjectType::empty()),
184        )]);
185        let schema = Schema::new("Test", SchemaType::Object(obj));
186        let config = GeneratorConfig::default();
187
188        let output = gen.generate(&schema, &config).unwrap();
189
190        assert!(output.contains("z.record(z.unknown())"));
191    }
192}