fork of hey-api/openapi-ts because I need some additional things
1import type { AnalysisContext, NodeName, Ref } from '@hey-api/codegen-core';
2import { ref } from '@hey-api/codegen-core';
3import type { MaybeArray } from '@hey-api/types';
4
5import { py } from '../../py-compiler';
6import type { MaybePyDsl } from '../base';
7import { PyDsl } from '../base';
8import type { DoExpr } from '../mixins/do';
9import { BlockPyDsl } from './block';
10
11const Mixed = PyDsl<py.TryStatement>;
12
13type ExceptType = string | MaybePyDsl<py.Expression>;
14
15interface ExceptEntry {
16 body: Array<DoExpr>;
17 name?: Ref<NodeName>;
18 types: Array<ExceptType>;
19}
20
21function exceptKey(types: Array<ExceptType>): string {
22 return types
23 .map((t) => (typeof t === 'string' ? t : '<<expr>>'))
24 .sort()
25 .join(',');
26}
27
28export class TryPyDsl extends Mixed {
29 readonly '~dsl' = 'TryPyDsl';
30
31 /**
32 * Ordered list of except clauses. We also keep a lookup map
33 * (`_exceptIndex`) keyed by the normalised type key so that
34 * repeated `.except()` calls with the same type set merge their
35 * body statements instead of creating duplicate clauses.
36 */
37 protected _excepts: Array<ExceptEntry> = [];
38 protected _exceptIndex: Map<string, number> = new Map();
39
40 protected _finally?: Array<DoExpr>;
41 protected _try?: Array<DoExpr>;
42
43 constructor(...tryBlock: Array<DoExpr>) {
44 super();
45 this.try(...tryBlock);
46 }
47
48 override analyze(ctx: AnalysisContext): void {
49 super.analyze(ctx);
50
51 if (this._try) {
52 ctx.pushScope();
53 try {
54 for (const stmt of this._try) ctx.analyze(stmt);
55 } finally {
56 ctx.popScope();
57 }
58 }
59
60 for (const entry of this._excepts) {
61 ctx.pushScope();
62 try {
63 ctx.analyze(entry.name);
64 for (const t of entry.types) ctx.analyze(t);
65 for (const stmt of entry.body) ctx.analyze(stmt);
66 } finally {
67 ctx.popScope();
68 }
69 }
70
71 if (this._finally) {
72 ctx.pushScope();
73 try {
74 for (const stmt of this._finally) ctx.analyze(stmt);
75 } finally {
76 ctx.popScope();
77 }
78 }
79 }
80
81 /** Returns true when all required builder calls are present. */
82 get isValid(): boolean {
83 return !this.missingRequiredCalls().length;
84 }
85
86 /**
87 * Add (or merge into) an except clause.
88 *
89 * ```ts
90 * $.try(...)
91 * .except('ValueError', 'e', body1, body2) // except ValueError as e:
92 * .except(['TypeError', 'KeyError'], 'e', ...) // except (TypeError, KeyError) as e:
93 * .except('ValueError', moreBody) // merges into first clause
94 * ```
95 *
96 * @param types Single exception type or array of types.
97 * @param nameOrBody Either the `as` variable name (`NodeName`) or the
98 * first body expression. If it looks like a `NodeName` (string that
99 * is a valid Python identifier and is *not* a DSL node), it is treated
100 * as the name; pass body items after it.
101 * @param body Remaining body statements.
102 */
103 except(
104 types: MaybeArray<ExceptType>,
105 nameOrBody?: NodeName | DoExpr,
106 ...body: Array<DoExpr>
107 ): this {
108 const typeArr = Array.isArray(types) ? types : [types];
109 const key = exceptKey(typeArr);
110
111 let name: Ref<NodeName> | undefined;
112 let bodyItems: Array<DoExpr>;
113
114 // Disambiguate: if the second arg is a plain string that looks like
115 // an identifier (no dots, no spaces, not a DSL node) treat it as
116 // the `as` name. Otherwise it's the first body expression.
117 if (nameOrBody !== undefined && this._isNodeName(nameOrBody)) {
118 name = ref(nameOrBody as NodeName);
119 bodyItems = body;
120 } else if (nameOrBody !== undefined) {
121 bodyItems = [nameOrBody as DoExpr, ...body];
122 } else {
123 bodyItems = body;
124 }
125
126 const existing = this._exceptIndex.get(key);
127 if (existing !== undefined) {
128 const entry = this._excepts[existing]!;
129 entry.body.push(...bodyItems);
130 if (name !== undefined) entry.name = name;
131 } else {
132 this._exceptIndex.set(key, this._excepts.length);
133 this._excepts.push({ body: bodyItems, name, types: typeArr });
134 }
135
136 return this;
137 }
138
139 /** Add a bare `except:` clause (catches everything). */
140 exceptAll(...body: Array<DoExpr>): this {
141 const key = '';
142 const existing = this._exceptIndex.get(key);
143 if (existing !== undefined) {
144 this._excepts[existing]!.body.push(...body);
145 } else {
146 this._exceptIndex.set(key, this._excepts.length);
147 this._excepts.push({ body, types: [] });
148 }
149 return this;
150 }
151
152 finally(...items: Array<DoExpr>): this {
153 this._finally = items;
154 return this;
155 }
156
157 try(...items: Array<DoExpr>): this {
158 this._try = items;
159 return this;
160 }
161
162 override toAst() {
163 this.$validate();
164
165 const tryStatements = new BlockPyDsl(...this._try!).$do();
166
167 let exceptClauses: Array<py.ExceptClause> | undefined;
168 if (this._excepts.length) {
169 exceptClauses = this._excepts.map((entry) => {
170 const bodyStatements = new BlockPyDsl(...entry.body).$do();
171
172 let exceptionType: py.Expression | undefined;
173 if (entry.types.length === 1) {
174 exceptionType = this.$node(entry.types[0]!);
175 } else if (entry.types.length > 1) {
176 exceptionType = py.factory.createTupleExpression(entry.types.map((t) => this.$node(t)));
177 }
178
179 const exceptionName = entry.name
180 ? py.factory.createIdentifier(this.$name(entry.name) || String(entry.name['~ref']))
181 : undefined;
182
183 return py.factory.createExceptClause([...bodyStatements], exceptionType, exceptionName);
184 });
185 }
186
187 const finallyStatements = this._finally
188 ? [...new BlockPyDsl(...this._finally).$do()]
189 : undefined;
190
191 return py.factory.createTryStatement(
192 [...tryStatements],
193 exceptClauses,
194 undefined,
195 finallyStatements,
196 );
197 }
198
199 $validate(): asserts this is this & {
200 _try: Array<DoExpr>;
201 } {
202 const missing = this.missingRequiredCalls();
203 if (!missing.length) return;
204 throw new Error(`Try statement missing ${missing.join(' and ')}`);
205 }
206
207 private missingRequiredCalls(): ReadonlyArray<string> {
208 const missing: Array<string> = [];
209 if (!this._try || !this._try.length) missing.push('.try()');
210 return missing;
211 }
212
213 /**
214 * Heuristic: a value is a `NodeName` (intended as the `as` variable)
215 * if it is a plain string matching a Python identifier pattern, or a
216 * Symbol.
217 */
218 private _isNodeName(value: unknown): boolean {
219 if (typeof value === 'string') {
220 return /^[A-Za-z_]\w*$/.test(value);
221 }
222 // Symbols from codegen-core have `~brand`
223 if (value && typeof value === 'object' && '~brand' in value) {
224 return true;
225 }
226 return false;
227 }
228}