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