import type { SymbolMeta } from '@hey-api/codegen-core'; import { fromRef, ref } from '@hey-api/codegen-core'; import type { IR, SchemaWithType } from '@hey-api/shared'; import { deduplicateSchema, pathToJsonPointer } from '@hey-api/shared'; // import { $ } from '../../../py-dsl'; import type { Ast, IrSchemaToAstOptions } from '../shared/types'; import type { PydanticPlugin } from '../types'; import { createProcessor } from './processor'; import { irSchemaWithTypeToAst } from './toAst'; export function irSchemaToAst({ optional, plugin, schema, schemaExtractor, state, }: IrSchemaToAstOptions & { optional?: boolean; schema: IR.SchemaObject; }): Ast { if (schemaExtractor && !schema.$ref) { const extracted = schemaExtractor({ meta: { resource: 'definition', resourceId: pathToJsonPointer(fromRef(state.path)), }, naming: plugin.config.definitions, path: fromRef(state.path), plugin, schema, }); if (extracted !== schema) schema = extracted; } if (schema.$ref) { const query: SymbolMeta = { category: 'schema', resource: 'definition', resourceId: schema.$ref, tool: 'pydantic', }; const refSymbol = plugin.referenceSymbol(query); const refName = typeof refSymbol === 'string' ? refSymbol : refSymbol.name; return { // expression: $.expr(refName), fieldConstraints: optional ? { default: null } : undefined, hasLazyExpression: !plugin.isSymbolRegistered(query), models: [], // pipes: [], typeAnnotation: refName, }; } if (schema.type) { const typeAst = irSchemaWithTypeToAst({ plugin, schema: schema as SchemaWithType, state, }); const constraints: Record = {}; if (optional) { constraints.default = null; } if (schema.default !== undefined) { constraints.default = schema.default; } if (schema.description) { constraints.description = schema.description; } return { ...typeAst, fieldConstraints: { ...typeAst.fieldConstraints, ...constraints }, // pipes: [], }; } if (schema.items) { schema = deduplicateSchema({ schema }); if (schema.items) { const itemsAnnotations: string[] = []; const itemsConstraints: Record[] = []; for (const item of schema.items) { const itemAst = irSchemaToAst({ plugin, schema: item, state: { ...state, path: ref([...fromRef(state.path), 'items']), }, }); itemsAnnotations.push(itemAst.typeAnnotation); if (itemAst.fieldConstraints) { itemsConstraints.push(itemAst.fieldConstraints); } } const unionType = itemsAnnotations.join(' | '); return { // expression: $.expr(`list[${unionType}]`), fieldConstraints: itemsConstraints.length > 0 ? itemsConstraints[0] : undefined, hasLazyExpression: false, models: [], // pipes: [], typeAnnotation: `list[${unionType}]`, }; } } return { // expression: $.expr('Any'), hasLazyExpression: false, models: [], // pipes: [], typeAnnotation: 'Any', }; } export const handlerV2: PydanticPlugin['Handler'] = ({ plugin }) => { plugin.symbol('Any', { external: 'typing', importKind: 'named', meta: { category: 'external', resource: 'typing.Any', }, }); plugin.symbol('BaseModel', { external: 'pydantic', importKind: 'named', meta: { category: 'external', resource: 'pydantic.BaseModel', }, }); plugin.symbol('ConfigDict', { external: 'pydantic', importKind: 'named', meta: { category: 'external', resource: 'pydantic.ConfigDict', }, }); plugin.symbol('Field', { external: 'pydantic', importKind: 'named', meta: { category: 'external', resource: 'pydantic.Field', }, }); plugin.symbol('Literal', { external: 'typing', importKind: 'named', meta: { category: 'external', resource: 'typing.Literal', }, }); plugin.symbol('Optional', { external: 'typing', importKind: 'named', meta: { category: 'external', resource: 'typing.Optional', }, }); const processor = createProcessor(plugin); plugin.forEach('operation', 'parameter', 'requestBody', 'schema', 'webhook', (event) => { switch (event.type) { case 'parameter': processor.process({ meta: { resource: 'definition', resourceId: pathToJsonPointer(event._path), }, naming: plugin.config.definitions, path: event._path, plugin, schema: event.parameter.schema, tags: event.tags, }); break; case 'requestBody': processor.process({ meta: { resource: 'definition', resourceId: pathToJsonPointer(event._path), }, naming: plugin.config.definitions, path: event._path, plugin, schema: event.requestBody.schema, tags: event.tags, }); break; case 'schema': processor.process({ meta: { resource: 'definition', resourceId: pathToJsonPointer(event._path), }, naming: plugin.config.definitions, path: event._path, plugin, schema: event.schema, tags: event.tags, }); break; } }); };