3x-ui/tools/openapigen/emit_zod.go

package main

import (
	"fmt"
	"io"
	"sort"
	"strings"
)

func emitZod(w io.Writer, schemas []Schema, aliases []Alias) error {
	if _, err := fmt.Fprintln(w, zodHeader); err != nil {
		return err
	}

	for _, a := range sortAliases(aliases) {
		if _, err := fmt.Fprintf(w, "export const %sSchema = %s;\n", a.Name, zodTypeExpr(a.Underlying)); err != nil {
			return err
		}
		if _, err := fmt.Fprintf(w, "export type %s = z.infer<typeof %sSchema>;\n\n", a.Name, a.Name); err != nil {
			return err
		}
	}

	for _, s := range sortSchemas(schemas) {
		if _, err := fmt.Fprintf(w, "export const %sSchema = z.object({\n", s.Name); err != nil {
			return err
		}
		fields := append([]Field(nil), s.Fields...)
		sort.SliceStable(fields, func(i, j int) bool { return fields[i].JSONName < fields[j].JSONName })
		for _, f := range fields {
			line := fmt.Sprintf("  %s: %s,\n", quoteIfNeeded(f.JSONName), zodExpr(f))
			if _, err := fmt.Fprint(w, line); err != nil {
				return err
			}
		}
		if _, err := fmt.Fprintln(w, "});"); err != nil {
			return err
		}
		if _, err := fmt.Fprintf(w, "export type %s = z.infer<typeof %sSchema>;\n\n", s.Name, s.Name); err != nil {
			return err
		}
	}
	return nil
}

func zodExpr(f Field) string {
	expr := zodTypeExpr(f.Type)
	expr = applyZodValidations(expr, f.Type, f.Validate)
	if f.Optional {
		expr += ".optional()"
	}
	return expr
}

func zodTypeExpr(t TypeRef) string {
	switch t.Kind {
	case KindString:
		return "z.string()"
	case KindBool:
		return "z.boolean()"
	case KindInt:
		return "z.number().int()"
	case KindNumber:
		return "z.number()"
	case KindAny, KindUnknown:
		return "z.unknown()"
	case KindRaw:
		return "z.unknown()"
	case KindArray:
		return "z.array(" + zodTypeExpr(*t.Element) + ")"
	case KindMap:
		return "z.record(" + zodTypeExpr(*t.Key) + ", " + zodTypeExpr(*t.Value) + ")"
	case KindRef:
		if t.Name == "nullable" {
			return zodTypeExpr(*t.Inner) + ".nullable()"
		}
		return "z.lazy(() => " + t.Name + "Schema)"
	}
	return "z.unknown()"
}

func applyZodValidations(expr string, t TypeRef, rules []ValidateRule) string {
	for _, r := range rules {
		switch r.Name {
		case "required":
			continue
		case "omitempty":
			continue
		case "gte":
			if t.Kind == KindInt || t.Kind == KindNumber {
				expr += fmt.Sprintf(".min(%s)", r.Param)
			}
		case "lte":
			if t.Kind == KindInt || t.Kind == KindNumber {
				expr += fmt.Sprintf(".max(%s)", r.Param)
			}
		case "gt":
			if t.Kind == KindInt || t.Kind == KindNumber {
				expr += fmt.Sprintf(".gt(%s)", r.Param)
			}
		case "lt":
			if t.Kind == KindInt || t.Kind == KindNumber {
				expr += fmt.Sprintf(".lt(%s)", r.Param)
			}
		case "min":
			switch t.Kind {
			case KindString:
				expr += fmt.Sprintf(".min(%s)", r.Param)
			case KindInt, KindNumber:
				expr += fmt.Sprintf(".min(%s)", r.Param)
			}
		case "max":
			switch t.Kind {
			case KindString:
				expr += fmt.Sprintf(".max(%s)", r.Param)
			case KindInt, KindNumber:
				expr += fmt.Sprintf(".max(%s)", r.Param)
			}
		case "url":
			expr += ".url()"
		case "email":
			expr += ".email()"
		case "oneof":
			values := strings.Fields(r.Param)
			quoted := make([]string, 0, len(values))
			for _, v := range values {
				quoted = append(quoted, fmt.Sprintf("'%s'", v))
			}
			expr = fmt.Sprintf("z.enum([%s])", strings.Join(quoted, ", "))
		}
	}
	return expr
}

func quoteIfNeeded(name string) string {
	if name == "" {
		return "''"
	}
	for i, r := range name {
		if r >= 'a' && r <= 'z' {
			continue
		}
		if r >= 'A' && r <= 'Z' {
			continue
		}
		if r == '_' || r == '$' {
			continue
		}
		if i > 0 && r >= '0' && r <= '9' {
			continue
		}
		return "'" + name + "'"
	}
	return name
}

const zodHeader = `// Code generated by tools/openapigen. DO NOT EDIT.
import { z } from 'zod';`