fork of hey-api/openapi-ts because I need some additional things
1import type { Symbol } from '@hey-api/codegen-core';
2import { buildSymbolIn, pathToName } from '@hey-api/shared';
3
4import { $ } from '../../../py-dsl';
5import type { PydanticPlugin } from '../types';
6import { identifiers } from '../v2/constants';
7import { createFieldCall } from './field';
8import type { ProcessorContext } from './processor';
9import type { PydanticField, PydanticFinal } from './types';
10
11export function exportAst({
12 final,
13 meta,
14 naming,
15 namingAnchor,
16 path,
17 plugin,
18 schema,
19 tags,
20}: ProcessorContext & {
21 final: PydanticFinal;
22}): void {
23 const name = pathToName(path, { anchor: namingAnchor });
24 const symbol = plugin.registerSymbol(
25 buildSymbolIn({
26 meta: {
27 category: 'schema',
28 path,
29 tags,
30 tool: 'pydantic',
31 ...meta,
32 },
33 name,
34 naming,
35 plugin,
36 schema,
37 }),
38 );
39
40 if (final.enumMembers) {
41 exportEnumClass({ final, plugin, symbol });
42 } else if (final.fields) {
43 exportClass({ final, plugin, symbol });
44 } else {
45 exportTypeAlias({ final, plugin, symbol });
46 }
47}
48
49function exportClass({
50 final,
51 plugin,
52 symbol,
53}: {
54 final: PydanticFinal;
55 plugin: PydanticPlugin['Instance'];
56 symbol: Symbol;
57}): void {
58 const baseModel = plugin.external('pydantic.BaseModel');
59 const classDef = $.class(symbol).extends(baseModel);
60
61 if (plugin.config.strict) {
62 const configDict = plugin.external('pydantic.ConfigDict');
63 classDef.do(
64 $.var(identifiers.model_config).assign($(configDict).call($.kwarg('extra', 'forbid'))),
65 );
66 }
67
68 for (const field of final.fields!) {
69 const fieldStatement = createFieldStatement(field, plugin);
70 classDef.do(fieldStatement);
71 }
72
73 plugin.node(classDef);
74}
75
76function exportEnumClass({
77 final,
78 plugin,
79 symbol,
80}: {
81 final: PydanticFinal;
82 plugin: PydanticPlugin['Instance'];
83 symbol: Symbol;
84}): void {
85 const members = final.enumMembers ?? [];
86 const hasStrings = members.some((m) => typeof m.value === 'string');
87 const hasNumbers = members.some((m) => typeof m.value === 'number');
88
89 const enumSymbol = plugin.external('enum.Enum');
90 const classDef = $.class(symbol).extends(enumSymbol);
91
92 if (hasStrings && !hasNumbers) {
93 classDef.extends('str');
94 } else if (!hasStrings && hasNumbers) {
95 classDef.extends('int');
96 }
97
98 for (const member of final.enumMembers ?? []) {
99 classDef.do($.var(member.name).assign($.literal(member.value)));
100 }
101
102 plugin.node(classDef);
103}
104
105function createFieldStatement(
106 field: PydanticField,
107 plugin: PydanticPlugin['Instance'],
108): ReturnType<typeof $.var> {
109 const fieldSymbol = field.name;
110 const varStatement = $.var(fieldSymbol).$if(field.type, (v, a) => v.type(a));
111
112 const originalName = field.originalName ?? fieldSymbol.name;
113 const needsAlias = field.originalName !== undefined && fieldSymbol.name !== originalName;
114
115 const constraints = {
116 ...field.fieldConstraints,
117 ...(needsAlias && !field.fieldConstraints?.alias && { alias: originalName }),
118 };
119
120 if (Object.keys(constraints).length) {
121 const fieldCall = createFieldCall(constraints, plugin, {
122 required: !field.isOptional,
123 });
124 return varStatement.assign(fieldCall);
125 }
126
127 if (field.isOptional) {
128 return varStatement.assign('None');
129 }
130
131 return varStatement;
132}
133
134function exportTypeAlias({
135 final,
136 plugin,
137 symbol,
138}: {
139 final: PydanticFinal;
140 plugin: PydanticPlugin['Instance'];
141 symbol: Symbol;
142}): void {
143 const typeAlias = plugin.external('typing.TypeAlias');
144 const statement = $.var(symbol)
145 .type(typeAlias)
146 .assign(final.type ?? plugin.external('typing.Any'));
147 plugin.node(statement);
148}