fork of hey-api/openapi-ts because I need some additional things
1import type { SchemaWithType } from '@hey-api/shared';
2import { toCase } from '@hey-api/shared';
3
4import { $ } from '../../../../py-dsl';
5import type { EnumResolverContext } from '../../resolvers';
6import type { PydanticFinal, PydanticType } from '../../shared/types';
7import type { PydanticPlugin } from '../../types';
8
9export interface EnumToTypeResult extends PydanticType {
10 enumMembers: Required<PydanticFinal>['enumMembers'];
11 isNullable: boolean;
12}
13
14// TODO: replace with casing utils
15function toEnumMemberName(value: string | number): string {
16 if (typeof value === 'number') {
17 // For numbers, prefix with underscore if starts with digit
18 return `VALUE_${value}`.replace(/-/g, '_NEG_').replace(/\./g, '_DOT_');
19 }
20
21 return toCase(value, 'SCREAMING_SNAKE_CASE');
22}
23
24function itemsNode(ctx: EnumResolverContext) {
25 const { plugin, schema } = ctx;
26 const enumMembers: Required<PydanticFinal>['enumMembers'] = [];
27 let isNullable = false;
28
29 for (const item of schema.items ?? []) {
30 if (item.type === 'null' || item.const === null) {
31 isNullable = true;
32 continue;
33 }
34
35 if (
36 (item.type === 'string' && typeof item.const === 'string') ||
37 ((item.type === 'integer' || item.type === 'number') && typeof item.const === 'number')
38 ) {
39 enumMembers.push({
40 name: plugin.symbol(toEnumMemberName(item.const)),
41 value: item.const,
42 });
43 }
44 }
45
46 return { enumMembers, isNullable };
47}
48
49function baseNode(ctx: EnumResolverContext): PydanticType {
50 const { plugin } = ctx;
51 const { enumMembers } = ctx.nodes.items(ctx);
52
53 if (!enumMembers.length) {
54 return {
55 type: plugin.external('typing.Any'),
56 };
57 }
58
59 const mode = plugin.config.enums ?? 'enum';
60
61 if (mode === 'literal') {
62 if (!enumMembers.length) {
63 return {
64 type: plugin.external('typing.Any'),
65 };
66 }
67
68 const literal = plugin.external('typing.Literal');
69 const values = enumMembers.map((m) =>
70 // TODO: replace
71 typeof m.value === 'string' ? `"<<<<${m.value}"` : `<<<${m.value}`,
72 );
73
74 return {
75 type: $(literal).slice(...values),
76 };
77 }
78
79 return {};
80}
81
82function enumResolver(ctx: EnumResolverContext): PydanticType {
83 return ctx.nodes.base(ctx);
84}
85
86export function enumToType({
87 plugin,
88 schema,
89}: {
90 plugin: PydanticPlugin['Instance'];
91 schema: SchemaWithType<'enum'>;
92}): EnumToTypeResult {
93 const ctx: EnumResolverContext = {
94 $,
95 nodes: {
96 base: baseNode,
97 items: itemsNode,
98 },
99 plugin,
100 schema,
101 };
102
103 const resolver = plugin.config['~resolvers']?.enum;
104 const resolved = resolver?.(ctx) ?? enumResolver(ctx);
105
106 const { enumMembers, isNullable } = ctx.nodes.items(ctx);
107
108 return {
109 ...resolved,
110 enumMembers,
111 isNullable,
112 };
113}