fork of hey-api/openapi-ts because I need some additional things
1import type { PyNode } from './nodes/base';
2import { PyNodeKind } from './nodes/kinds';
3
4export interface PyPrinterOptions {
5 indentSize?: number;
6}
7
8const DEFAULT_INDENT_SIZE = 4;
9const PARAMS_MULTILINE_THRESHOLD = 3;
10
11export function createPrinter(options?: PyPrinterOptions) {
12 const indentSize = options?.indentSize ?? DEFAULT_INDENT_SIZE;
13
14 let indentLevel = 0;
15
16 function printComments(
17 parts: Array<string>,
18 lines: ReadonlyArray<string>,
19 indent?: boolean,
20 ): void {
21 if (indent) indentLevel += 1;
22 parts.push(...lines.map((line) => printLine(`# ${line}`)));
23 if (indent) indentLevel -= 1;
24 }
25
26 function printDocstring(docstring: string): Array<string> {
27 const lines = docstring.split('\n');
28 const parts: Array<string> = [];
29 if (lines.length === 1) {
30 parts.push(printLine(`"""${lines[0]}"""`), '');
31 } else {
32 parts.push(printLine(`"""`));
33 parts.push(...lines.map((line) => printLine(line)));
34 parts.push(printLine(`"""`), '');
35 }
36 return parts;
37 }
38
39 function printLine(line: string): string {
40 if (line === '') return '';
41 return ' '.repeat(indentLevel * indentSize) + line;
42 }
43
44 function printNode(node: PyNode): string {
45 const parts: Array<string> = [];
46
47 if (node.leadingComments) {
48 printComments(parts, node.leadingComments);
49 }
50
51 let indentTrailingComments = false;
52
53 switch (node.kind) {
54 case PyNodeKind.Assignment: {
55 const target = printNode(node.target);
56 if (node.type) {
57 const type = printNode(node.type);
58 if (node.value) {
59 parts.push(printLine(`${target}: ${type} = ${printNode(node.value)}`));
60 } else {
61 parts.push(printLine(`${target}: ${type}`));
62 }
63 } else {
64 parts.push(printLine(`${target} = ${printNode(node.value!)}`));
65 }
66 break;
67 }
68
69 case PyNodeKind.AsyncExpression:
70 parts.push(`async ${printNode(node.expression)}`);
71 break;
72
73 case PyNodeKind.AugmentedAssignment:
74 parts.push(
75 printLine(`${printNode(node.target)} ${node.operator} ${printNode(node.value)}`),
76 );
77 break;
78
79 case PyNodeKind.AwaitExpression:
80 parts.push(`await ${printNode(node.expression)}`);
81 break;
82
83 case PyNodeKind.BinaryExpression:
84 parts.push(`${printNode(node.left)} ${node.operator} ${printNode(node.right)}`);
85 break;
86
87 case PyNodeKind.Block:
88 indentLevel += 1;
89 if (node.statements.length) {
90 parts.push(...node.statements.map(printNode));
91 } else {
92 parts.push(printLine('pass'));
93 }
94 indentLevel -= 1;
95 break;
96
97 case PyNodeKind.BreakStatement:
98 parts.push(printLine('break'));
99 break;
100
101 case PyNodeKind.CallExpression:
102 parts.push(`${printNode(node.callee)}(${node.args.map(printNode).join(', ')})`);
103 break;
104
105 case PyNodeKind.ClassDeclaration: {
106 indentTrailingComments = true;
107 if (node.decorators) {
108 parts.push(...node.decorators.map((decorator) => printLine(`@${printNode(decorator)}`)));
109 }
110 const bases = node.baseClasses?.length
111 ? `(${node.baseClasses.map(printNode).join(', ')})`
112 : '';
113 parts.push(printLine(`class ${node.name}${bases}:`));
114 if (node.docstring) {
115 indentLevel += 1;
116 parts.push(...printDocstring(node.docstring));
117 indentLevel -= 1;
118 }
119 parts.push(printNode(node.body));
120 break;
121 }
122
123 case PyNodeKind.Comment:
124 parts.push(printLine(`# ${node.text}`));
125 break;
126
127 case PyNodeKind.ContinueStatement:
128 parts.push(printLine('continue'));
129 break;
130
131 case PyNodeKind.DictComprehension: {
132 const asyncPrefix = node.isAsync ? 'async ' : '';
133 const children: Array<string> = [
134 `${printNode(node.key)}: ${printNode(node.value)} ${asyncPrefix}for ${printNode(node.target)} in ${printNode(node.iterable)}`,
135 ];
136 if (node.ifs) {
137 for (const condition of node.ifs) {
138 children.push(`if ${printNode(condition)}`);
139 }
140 }
141 parts.push(`{${children.join(' ')}}`);
142 break;
143 }
144
145 case PyNodeKind.DictExpression: {
146 const entries = node.entries
147 .map(({ key, value }) => `${printNode(key)}: ${printNode(value)}`)
148 .join(', ');
149 parts.push(`{${entries}}`);
150 break;
151 }
152
153 case PyNodeKind.EmptyStatement:
154 parts.push('');
155 break;
156
157 case PyNodeKind.ExpressionStatement:
158 parts.push(printLine(printNode(node.expression)));
159 break;
160
161 case PyNodeKind.ForStatement:
162 parts.push(printLine(`for ${printNode(node.target)} in ${printNode(node.iterable)}:`));
163 parts.push(printNode(node.body));
164 if (node.elseBlock) {
165 parts.push(`${printLine('else:')}`);
166 parts.push(`${printNode(node.elseBlock)}`);
167 }
168 break;
169
170 case PyNodeKind.FStringExpression: {
171 const children = node.parts.map((part) =>
172 typeof part === 'string' ? part : `{${printNode(part)}}`,
173 );
174 parts.push(`f"${children.join('')}"`);
175 break;
176 }
177
178 case PyNodeKind.FunctionDeclaration: {
179 if (node.decorators) {
180 parts.push(...node.decorators.map((decorator) => printLine(`@${printNode(decorator)}`)));
181 }
182 const modifiers = node.modifiers?.map(printNode).join(' ') ?? '';
183 const defPrefix = modifiers ? `${modifiers} def` : 'def';
184 const formatParameter = (parameter: (typeof node.parameters)[number]): string => {
185 const children: Array<string> = [parameter.name];
186 if (parameter.type) children.push(`: ${printNode(parameter.type)}`);
187 if (parameter.defaultValue) children.push(` = ${printNode(parameter.defaultValue)}`);
188 return children.join('');
189 };
190 const returnAnnotation = node.returnType ? ` -> ${printNode(node.returnType)}` : '';
191 const preParams = `${defPrefix} ${node.name}(`;
192 const postParams = `)${returnAnnotation}:`;
193
194 if (node.parameters.length > PARAMS_MULTILINE_THRESHOLD) {
195 parts.push(printLine(preParams));
196 indentLevel += 1;
197 const params = node.parameters.map((parameter) =>
198 printLine(`${formatParameter(parameter)},`),
199 );
200 parts.push(...params);
201 indentLevel -= 1;
202 parts.push(printLine(postParams));
203 } else {
204 const parameters = node.parameters.map((parameter) => formatParameter(parameter));
205 const params = parameters.join(', ');
206 parts.push(printLine(`${preParams}${params}${postParams}`));
207 }
208
209 if (node.docstring) {
210 indentLevel += 1;
211 parts.push(...printDocstring(node.docstring));
212 indentLevel -= 1;
213 }
214 parts.push(printNode(node.body));
215 break;
216 }
217
218 case PyNodeKind.GeneratorExpression: {
219 const asyncPrefix = node.isAsync ? 'async ' : '';
220 const children: Array<string> = [
221 `${printNode(node.element)} ${asyncPrefix}for ${printNode(node.target)} in ${printNode(node.iterable)}`,
222 ];
223 if (node.ifs) {
224 for (const condition of node.ifs) {
225 children.push(`if ${printNode(condition)}`);
226 }
227 }
228 parts.push(`(${children.join(' ')})`);
229 break;
230 }
231
232 case PyNodeKind.Identifier:
233 parts.push(node.name);
234 break;
235
236 case PyNodeKind.IfStatement:
237 parts.push(printLine(`if ${printNode(node.condition)}:`));
238 parts.push(`${printNode(node.thenBlock)}`);
239 if (node.elseBlock) {
240 parts.push(`${printLine('else:')}`);
241 parts.push(`${printNode(node.elseBlock)}`);
242 }
243 break;
244
245 case PyNodeKind.KeywordArgument:
246 parts.push(`${node.name}=${printNode(node.value)}`);
247 break;
248
249 case PyNodeKind.ImportStatement: {
250 const fromPrefix = node.isFrom ? `from ${node.module} ` : '';
251 if (fromPrefix) {
252 if (node.names && node.names.length) {
253 const imports = node.names
254 .map(({ alias, name }) => (alias ? `${name} as ${alias}` : name))
255 .join(', ');
256 parts.push(printLine(`${fromPrefix}import ${imports}`));
257 } else {
258 parts.push(printLine(`${fromPrefix}import *`));
259 }
260 } else {
261 if (node.names && node.names.length) {
262 const imports = node.names
263 .map(({ alias, name }) => (alias ? `${name} as ${alias}` : name))
264 .join(', ');
265 parts.push(printLine(`import ${imports}`));
266 } else {
267 parts.push(printLine(`import ${node.module}`));
268 }
269 }
270 break;
271 }
272
273 case PyNodeKind.LambdaExpression: {
274 const parameters = node.parameters.map((parameter) => {
275 const children: Array<string> = [parameter.name];
276 if (parameter.type) children.push(`: ${printNode(parameter.type)}`);
277 if (parameter.defaultValue) children.push(` = ${printNode(parameter.defaultValue)}`);
278 return children.join('');
279 });
280 parts.push(`lambda ${parameters.join(', ')}: ${printNode(node.expression)}`);
281 break;
282 }
283
284 case PyNodeKind.ListComprehension: {
285 const asyncPrefix = node.isAsync ? 'async ' : '';
286 const children: Array<string> = [
287 `${printNode(node.element)} ${asyncPrefix}for ${printNode(node.target)} in ${printNode(node.iterable)}`,
288 ];
289 if (node.ifs) {
290 for (const condition of node.ifs) {
291 children.push(`if ${printNode(condition)}`);
292 }
293 }
294 parts.push(`[${children.join(' ')}]`);
295 break;
296 }
297
298 case PyNodeKind.ListExpression:
299 parts.push(`[${node.elements.map(printNode).join(', ')}]`);
300 break;
301
302 case PyNodeKind.Literal:
303 if (typeof node.value === 'string') {
304 parts.push(`"${node.value}"`);
305 } else if (typeof node.value === 'boolean') {
306 parts.push(node.value ? 'True' : 'False');
307 } else if (node.value === null) {
308 parts.push('None');
309 } else {
310 parts.push(String(node.value));
311 }
312 break;
313
314 case PyNodeKind.MemberExpression:
315 parts.push(`${printNode(node.object)}.${printNode(node.member)}`);
316 break;
317
318 case PyNodeKind.RaiseStatement:
319 if (node.expression) {
320 parts.push(printLine(`raise ${printNode(node.expression)}`));
321 } else {
322 parts.push(printLine('raise'));
323 }
324 break;
325
326 case PyNodeKind.ReturnStatement:
327 if (node.expression) {
328 parts.push(printLine(`return ${printNode(node.expression)}`));
329 } else {
330 parts.push(printLine('return'));
331 }
332 break;
333
334 case PyNodeKind.SetComprehension: {
335 const asyncPrefix = node.isAsync ? 'async ' : '';
336 const children: Array<string> = [
337 `${printNode(node.element)} ${asyncPrefix}for ${printNode(node.target)} in ${printNode(node.iterable)}`,
338 ];
339 if (node.ifs) {
340 for (const condition of node.ifs) {
341 children.push(`if ${printNode(condition)}`);
342 }
343 }
344 parts.push(`{${children.join(' ')}}`);
345 break;
346 }
347
348 case PyNodeKind.SetExpression: {
349 if (!node.elements.length) {
350 parts.push('set()');
351 } else {
352 parts.push(`{${node.elements.map(printNode).join(', ')}}`);
353 }
354 break;
355 }
356
357 case PyNodeKind.SourceFile:
358 if (node.docstring) {
359 parts.push(...printDocstring(node.docstring));
360 }
361 parts.push(...node.statements.map(printNode));
362 break;
363
364 case PyNodeKind.SubscriptExpression:
365 parts.push(`${printNode(node.value)}[${printNode(node.slice)}]`);
366 break;
367
368 case PyNodeKind.SubscriptSlice:
369 parts.push(node.elements.map(printNode).join(', '));
370 break;
371
372 case PyNodeKind.TryStatement: {
373 parts.push(printLine('try:'), printNode(node.tryBlock));
374 if (node.exceptClauses) {
375 for (const clause of node.exceptClauses) {
376 const type = clause.exceptionType ? ` ${printNode(clause.exceptionType)}` : '';
377 const name = clause.exceptionName ? ` as ${printNode(clause.exceptionName)}` : '';
378 parts.push(printLine(`except${type}${name}:`), printNode(clause.block));
379 }
380 }
381 if (node.elseBlock) {
382 parts.push(printLine(`else:`), printNode(node.elseBlock));
383 }
384 if (node.finallyBlock) {
385 parts.push(printLine(`finally:`), printNode(node.finallyBlock));
386 }
387 break;
388 }
389
390 case PyNodeKind.TupleExpression: {
391 // Single-element tuple needs trailing comma
392 const trailingComma = node.elements.length === 1 ? ',' : '';
393 parts.push(`(${node.elements.map(printNode).join(', ')}${trailingComma})`);
394 break;
395 }
396
397 case PyNodeKind.WhileStatement: {
398 parts.push(printLine(`while ${printNode(node.condition)}:`));
399 parts.push(printNode(node.body));
400 if (node.elseBlock) {
401 parts.push(`${printLine('else:')}`);
402 parts.push(`${printNode(node.elseBlock)}`);
403 }
404 break;
405 }
406
407 case PyNodeKind.WithStatement: {
408 const modifiers = node.modifiers?.map(printNode).join(' ') ?? '';
409 const withPrefix = modifiers ? `${modifiers} with` : 'with';
410 const items = node.items
411 .map((item) =>
412 item.alias
413 ? `${printNode(item.contextExpr)} as ${printNode(item.alias)}`
414 : printNode(item.contextExpr),
415 )
416 .join(', ');
417 parts.push(printLine(`${withPrefix} ${items}:`));
418 parts.push(printNode(node.body));
419 break;
420 }
421
422 case PyNodeKind.YieldExpression:
423 if (node.value) {
424 parts.push(`yield ${printNode(node.value)}`);
425 } else {
426 parts.push('yield');
427 }
428 break;
429
430 case PyNodeKind.YieldFromExpression:
431 parts.push(`yield from ${printNode(node.expression)}`);
432 break;
433
434 default:
435 throw new Error(`Unsupported node kind: ${(node as { kind: string }).kind}`);
436 }
437
438 if (node.trailingComments) {
439 printComments(parts, node.trailingComments, indentTrailingComments);
440 }
441
442 return parts.join('\n');
443 }
444
445 function printFile(node: PyNode): string {
446 const parts: Array<string> = [printNode(node), ''];
447 return parts.join('\n');
448 }
449
450 return {
451 printFile,
452 };
453}
454
455export function printAst(node: PyNode): string {
456 return JSON.stringify(node, null, 2);
457}