zod-to-x
Version:
Multi language types generation from Zod schemas.
435 lines (434 loc) • 17.3 kB
JavaScript
"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.Zod2Py = void 0;
const case_1 = __importDefault(require("case"));
const core_1 = require("../../core");
const options_1 = require("./options");
const libs_1 = require("./libs");
class Zod2Py extends core_1.Zod2X {
constructor(opt = {}) {
super(Object.assign(Object.assign({}, options_1.defaultOpts), opt));
this.commentKey = "#";
this.baseSchemaAdded = false;
this.typeVars = new Set();
this.pendingTypeVars = new Set();
this.getAnyType = () => {
this.imports.add(this.lib.anyType);
return "Any";
};
this.getBooleanType = () => "bool";
this.getDateType = () => {
this.imports.add(this.lib.datetimeType);
return "datetime";
};
/** Ex: Set[TypeA] */
this.getSetType = (itemType) => {
this.imports.add(this.lib.setType);
return `Set[${itemType}]`;
};
this.getStringType = () => "str";
/** Ex: Tuple[TypeA, TypeB] */
this.getTupleType = (itemsType) => {
this.imports.add(this.lib.tupleType);
return `Tuple[${itemsType.join(", ")}]`;
};
/** Ex: Union[TypeA, TypeB] */
this.getUnionType = (itemsType) => {
this.imports.add(this.lib.unionType);
return `Union[${itemsType.join(", ")}]`;
};
/** Ex: TypeA & TypeB -> intersection handling */
this.getIntersectionType = (itemsType) => {
// Python doesn't have intersection types, we'll create a new class
return itemsType.join(" & "); // This will be handled in transpileIntersection
};
/** Ex: int or float depending on isInt flag */
this.getNumberType = (isInt, range) => {
return isInt ? "int" : "float";
};
this.lib = (0, libs_1.getLibs)();
}
runAfter() {
if (!this.baseSchemaAdded) {
this._flushPendingTypeVars(true);
}
this._consolidateImports();
}
runBefore() { }
/**
* Adds BaseSchema class definition if not already added.
* This is the base class for all Pydantic models with shared configuration.
*/
_addBaseSchema() {
if (this.baseSchemaAdded)
return;
this.baseSchemaAdded = true;
this.imports.add(this.lib.baseModel);
if (this.opt.keepKeys !== true) {
this.imports.add(this.lib.aliasGenerator);
}
this.push0("class BaseSchema(BaseModel):");
this.push1("model_config = ConfigDict(");
if (this.opt.keepKeys !== true) {
this.push2("alias_generator=to_camel,");
this.push2("serialize_by_alias=True,");
this.push2("populate_by_name=True,");
}
this.push2("use_enum_values=True");
this.push1(")");
this.push0("");
this._flushPendingTypeVars(false);
}
_flushPendingTypeVars(prepend) {
if (this.pendingTypeVars.size === 0)
return;
const pending = Array.from(this.pendingTypeVars);
const lines = pending.map((typeVar) => `${typeVar} = TypeVar('${typeVar}')`);
if (prepend) {
this.output = [...lines, "", ...this.output];
}
else {
lines.forEach((line) => this.push0(line));
this.push0("");
}
pending.forEach((typeVar) => this.typeVars.add(typeVar));
this.pendingTypeVars.clear();
}
/**
* Declares TypeVars that haven't been declared yet.
* Adds them right before their first usage.
* Ex: T = TypeVar('T')
*/
_declareNewTypeVars(templates) {
const newTypeVars = Array.from(templates).filter((t) => !this.typeVars.has(t) && !this.pendingTypeVars.has(t));
if (newTypeVars.length === 0)
return;
this.imports.add(this.lib.typeVarType);
if (!this.baseSchemaAdded) {
newTypeVars.forEach((typeVar) => this.pendingTypeVars.add(typeVar));
return;
}
newTypeVars.forEach((typeVar) => {
this.push0(`${typeVar} = TypeVar('${typeVar}')`);
this.typeVars.add(typeVar);
});
this.push0("");
}
/**
* Consolidates multiline imports from the same module and sorts them alphabetically.
* Modifies this.imports Set to contain consolidated import statements.
*/
_consolidateImports() {
const importGroups = new Map();
const simpleImports = new Set();
// Group imports by module
this.imports.forEach((importStatement) => {
if (importStatement.startsWith("from ") && importStatement.includes(" import ")) {
// Parse "from module import item" format
const match = importStatement.match(/^from\s+(\S+)\s+import\s+(.+)$/);
if (match) {
const module = match[1];
const items = match[2].split(",").map((item) => item.trim());
if (!importGroups.has(module)) {
importGroups.set(module, new Set());
}
items.forEach((item) => {
if (item) {
importGroups.get(module).add(item);
}
});
}
}
else if (importStatement.startsWith("import ")) {
// Simple import statement
simpleImports.add(importStatement);
}
});
// Clear existing imports
this.imports.clear();
// Add consolidated "from" imports back to this.imports (sorted by module)
const sortedModules = Array.from(importGroups.keys()).sort();
sortedModules.forEach((module) => {
const items = Array.from(importGroups.get(module)).sort();
if (items.length === 1) {
this.imports.add(`from ${module} import ${items[0]}`);
}
else {
this.imports.add(`from ${module} import ${items.join(", ")}`);
}
});
// Add simple imports back to this.imports (sorted)
const sortedSimpleImports = Array.from(simpleImports).sort();
sortedSimpleImports.forEach((imp) => {
this.imports.add(imp);
});
}
addImportFromFile(filename, namespace) {
const moduleName = filename.endsWith(".py") ? filename.slice(0, -3) : filename;
return `import ${moduleName} as ${namespace}`;
}
getTypeFromExternalNamespace(namespace, typeName) {
return `${namespace}.${typeName}`;
}
addExtendedType(name, parentNamespace, aliasOf, opt) {
var _a;
const extendedType = (opt === null || opt === void 0 ? void 0 : opt.isInternal)
? aliasOf
: this.getTypeFromExternalNamespace(parentNamespace, aliasOf);
const templates = (_a = opt === null || opt === void 0 ? void 0 : opt.templates) !== null && _a !== void 0 ? _a : "";
if (opt === null || opt === void 0 ? void 0 : opt.isClass) {
// For classes (ASTObject, ASTIntersection), use inheritance with templates
this.push0(`class ${name}(${extendedType}${templates}): ...\n`);
}
else {
// For type aliases (primitives, unions, etc.), use TypeAlias
this.imports.add(this.lib.typeAliasType);
this.push0(`${name}: TypeAlias = ${extendedType}${templates}\n`);
}
}
getGenericTemplatesTranslation(data) {
if ((data instanceof core_1.ASTObject || data instanceof core_1.ASTDefinition) &&
data.templatesTranslation.length > 0) {
return ("[" +
data.templatesTranslation
.map((t) => {
if (this.isExternalTypeImport(t)) {
this.addExternalTypeImport(t);
return this.getTypeFromExternalNamespace(t.parentNamespace, t.aliasOf);
}
else {
return t.aliasOf;
}
})
.join(", ") +
"]");
}
}
/**
* Emits an alias/extension declaration early for layered references.
* It keeps concrete template translations and falls back to declared templates (e.g. [T])
* for aliases of generic templates.
*/
checkExtendedTypeInclusion(data, type) {
// Determine if the aliased type is a class (ASTObject or ASTIntersection with newObject)
const isClass = data instanceof core_1.ASTObject ||
(data instanceof core_1.ASTIntersection && data.newObject !== undefined);
const declaredTemplatesFallback = data instanceof core_1.ASTObject && data.templates.size > 0 && this.isExternalTypeImport(data)
? `[${[...data.templates].join(", ")}]`
: undefined;
const translatedTemplates = this.getGenericTemplatesTranslation(data);
const templates = translatedTemplates || declaredTemplatesFallback;
if (!translatedTemplates && data instanceof core_1.ASTObject && data.templates.size > 0) {
this._declareNewTypeVars(data.templates);
}
if (this.isExternalTypeImport(data)) {
if (data.aliasOf) {
this.addExtendedType(data.name, data.parentNamespace, data.aliasOf, {
type,
templates,
isClass,
});
this.addExternalTypeImport(data);
}
return true;
}
else if (data.aliasOf) {
this.addExtendedType(data.name, data.parentNamespace, data.aliasOf, {
type,
isInternal: true,
templates,
isClass,
});
return true;
}
return false;
}
/** Ex: List[List[TypeA]] */
getArrayType(arrayType, arrayDeep) {
this.imports.add(this.lib.listType);
let output = `List[${arrayType}]`;
for (let i = 0; i < arrayDeep - 1; i++) {
output = `List[${output}]`;
}
return output;
}
/** Ex: Literal["value"] or Literal[true] or EnumName.ENUM_VALUE */
getLiteralStringType(value, parentEnumNameKey) {
if (!parentEnumNameKey)
this.imports.add(this.lib.literalType);
return ("Literal[" +
(parentEnumNameKey
? `${parentEnumNameKey[0]}.${case_1.default.constant(case_1.default.snake(parentEnumNameKey[1]))}`
: typeof value === "boolean"
? case_1.default.capital(value.toString())
: `${isNaN(Number(value)) ? `"${value}"` : value}`) +
"]");
}
/** Ex: Dict[TypeA, TypeB] */
getMapType(keyType, valueType) {
this.imports.add(this.lib.dictType);
return `Dict[${keyType}, ${valueType}]`;
}
/** Ex: Dict[TypeA, TypeB] */
getRecordType(keyType, valueType) {
this.imports.add(this.lib.dictType);
return `Dict[${keyType}, ${valueType}]`;
}
transpileAliasedType(data) {
if (this.checkExtendedTypeInclusion(data, "alias")) {
return;
}
let extendedType = undefined;
this.addComment(data.description);
if (data instanceof core_1.ASTArray) {
extendedType = this.getAttributeType(data.item);
}
else {
extendedType = this.getAttributeType(data);
}
if (extendedType !== undefined) {
this.imports.add(this.lib.typeAliasType);
this.push0(`${data.name}: TypeAlias = ${extendedType}\n`);
}
}
/** Ex:
* class MyEnum(str, Enum):
* ITEM_KEY1 = "ItemValue1"
* ITEM_KEY2 = "ItemValue2"
*
* # Or for mixed types:
* class MyEnum(Enum):
* ITEM_KEY1 = 1
* ITEM_KEY2 = "ItemValue2"
*/
transpileEnum(data) {
if (this.checkExtendedTypeInclusion(data, "alias")) {
return;
}
this.imports.add(this.lib.enumType);
this.addComment(data.description);
// Check if all values are strings
const allStrings = data.values.every(([, value]) => typeof value === "string");
const enumParent = allStrings ? "(str, Enum)" : "(Enum)";
this.push0(`class ${data.name}${enumParent}:`);
data.values.forEach(([key, value]) => {
const keyName = case_1.default.constant(case_1.default.snake(key));
const enumValue = typeof value === "string" ? `"${value}"` : `${value}`;
this.push1(`${keyName} = ${enumValue}`);
});
this.push0("");
}
transpileIntersection(data) {
if (this.checkExtendedTypeInclusion(data)) {
return;
}
this._addBaseSchema();
this.addComment(data.description);
// Use multiple inheritance like C++ for intersections
const leftType = this.getAttributeType(data.left);
const rightType = this.getAttributeType(data.right);
this.push0(`class ${data.name}(${leftType}, ${rightType}): ...\n`);
}
transpileStruct(data) {
if (this.checkExtendedTypeInclusion(data)) {
return;
}
this._addBaseSchema();
if (data.templates.size > 0) {
this._declareNewTypeVars(data.templates);
}
this.addComment(data.description);
this._transpileStructAsClass(data);
}
transpileUnion(data) {
if (this.checkExtendedTypeInclusion(data, data.discriminantKey === undefined ? "union" : "d-union")) {
return;
}
// Python uses Union type aliases (like C++), not merged classes
this._addBaseSchema();
this.addComment(data.description);
const attributesTypes = data.options.map(this.getAttributeType.bind(this));
// For discriminated unions, use Annotated with Field discriminator
if (data.discriminantKey) {
this.imports.add(this.lib.annotatedType);
this.imports.add(this.lib.fieldImport);
this.push0(`${data.name} = Annotated[`);
this.push1(`${this.getUnionType(attributesTypes)},`);
this.push1(`Field(discriminator='${data.discriminantKey}')`);
this.push0(`]\n`);
}
else {
this.push0(`${data.name} = ${this.getUnionType(attributesTypes)}\n`);
}
this._createUnionWrapper(data.name);
}
/**
* Creates a wrapper class for a union type.
* Python/Pydantic needs this for proper serialization/deserialization of unions.
* Ex:
* class UnionItemWrapper(BaseSchema):
* data: UnionItem
*/
_createUnionWrapper(unionName) {
const wrapperName = `${unionName}Wrapper`;
this.push0(`class ${wrapperName}(BaseSchema):`);
this.push1(`data: ${unionName}`);
this.push0("");
}
/** Ex:
* class MyStruct(BaseSchema):
* att1: TypeA
* att2: Optional[TypeB] = None
*
* # Or with generics:
* class MyGenericStruct(BaseSchema, Generic[T]):
* data: T
* */
_transpileStructAsClass(data) {
// Handle generic templates
let baseClasses = "BaseSchema";
if (data.templates.size > 0) {
const templates = Array.from(data.templates);
this.imports.add(this.lib.genericType);
baseClasses += `, Generic[${templates.join(", ")}]`;
}
this.push0(`class ${data.name}(${baseClasses}):`);
const hasProperties = Object.keys(data.properties).length > 0;
if (!hasProperties) {
this.push1("pass");
this.push0("");
return;
}
// Generate fields
for (const [key, value] of Object.entries(data.properties)) {
const keyName = this.opt.keepKeys === true ? key : case_1.default.snake(key);
this._transpileMember(keyName, value);
}
this.push0("");
}
/** For Class attributes.
* Ex: attribute1: Optional[TypeA] = None */
_transpileMember(memberName, memberNode) {
const pythonType = this.getAttributeType(memberNode);
const isOptional = memberNode.isOptional || memberNode.isNullable;
let typeAnnotation;
if (isOptional) {
this.imports.add(this.lib.optionalType);
typeAnnotation = pythonType.startsWith("Optional[")
? pythonType
: `Optional[${pythonType}]`;
}
else {
typeAnnotation = pythonType;
}
if (memberNode.description && !memberNode.name && !this.isTranspilerable(memberNode)) {
this.addComment(memberNode.description, `\n${this.indent[1]}`);
}
const defaultValue = isOptional ? " = None" : "";
this.push1(`${memberName}: ${typeAnnotation}${defaultValue}`);
}
}
exports.Zod2Py = Zod2Py;