1#include "Python.h"
2#include "pycore_ast.h" // expr_ty
3#include <float.h> // DBL_MAX_10_EXP
4#include <stdbool.h>
5
6/* This limited unparser is used to convert annotations back to strings
7 * during compilation rather than being a full AST unparser.
8 * See ast.unparse for a full unparser (written in Python)
9 */
10
11static PyObject *_str_open_br;
12static PyObject *_str_dbl_open_br;
13static PyObject *_str_close_br;
14static PyObject *_str_dbl_close_br;
15static PyObject *_str_inf;
16static PyObject *_str_replace_inf;
17
18/* Forward declarations for recursion via helper functions. */
19static PyObject *
20expr_as_unicode(expr_ty e, int level);
21static int
22append_ast_expr(_PyUnicodeWriter *writer, expr_ty e, int level);
23static int
24append_joinedstr(_PyUnicodeWriter *writer, expr_ty e, bool is_format_spec);
25static int
26append_formattedvalue(_PyUnicodeWriter *writer, expr_ty e);
27static int
28append_ast_slice(_PyUnicodeWriter *writer, expr_ty e);
29
30static int
31append_charp(_PyUnicodeWriter *writer, const char *charp)
32{
33 return _PyUnicodeWriter_WriteASCIIString(writer, charp, -1);
34}
35
36#define APPEND_STR_FINISH(str) do { \
37 return append_charp(writer, (str)); \
38 } while (0)
39
40#define APPEND_STR(str) do { \
41 if (-1 == append_charp(writer, (str))) { \
42 return -1; \
43 } \
44 } while (0)
45
46#define APPEND_STR_IF(cond, str) do { \
47 if ((cond) && -1 == append_charp(writer, (str))) { \
48 return -1; \
49 } \
50 } while (0)
51
52#define APPEND_STR_IF_NOT_FIRST(str) do { \
53 APPEND_STR_IF(!first, (str)); \
54 first = false; \
55 } while (0)
56
57#define APPEND_EXPR(expr, pr) do { \
58 if (-1 == append_ast_expr(writer, (expr), (pr))) { \
59 return -1; \
60 } \
61 } while (0)
62
63#define APPEND(type, value) do { \
64 if (-1 == append_ast_ ## type(writer, (value))) { \
65 return -1; \
66 } \
67 } while (0)
68
69static int
70append_repr(_PyUnicodeWriter *writer, PyObject *obj)
71{
72 PyObject *repr = PyObject_Repr(obj);
73
74 if (!repr) {
75 return -1;
76 }
77
78 if ((PyFloat_CheckExact(obj) && Py_IS_INFINITY(PyFloat_AS_DOUBLE(obj))) ||
79 PyComplex_CheckExact(obj))
80 {
81 PyObject *new_repr = PyUnicode_Replace(
82 repr,
83 _str_inf,
84 _str_replace_inf,
85 -1
86 );
87 Py_DECREF(repr);
88 if (!new_repr) {
89 return -1;
90 }
91 repr = new_repr;
92 }
93 int ret = _PyUnicodeWriter_WriteStr(writer, repr);
94 Py_DECREF(repr);
95 return ret;
96}
97
98/* Priority levels */
99
100enum {
101 PR_TUPLE,
102 PR_TEST, /* 'if'-'else', 'lambda' */
103 PR_OR, /* 'or' */
104 PR_AND, /* 'and' */
105 PR_NOT, /* 'not' */
106 PR_CMP, /* '<', '>', '==', '>=', '<=', '!=',
107 'in', 'not in', 'is', 'is not' */
108 PR_EXPR,
109 PR_BOR = PR_EXPR, /* '|' */
110 PR_BXOR, /* '^' */
111 PR_BAND, /* '&' */
112 PR_SHIFT, /* '<<', '>>' */
113 PR_ARITH, /* '+', '-' */
114 PR_TERM, /* '*', '@', '/', '%', '//' */
115 PR_FACTOR, /* unary '+', '-', '~' */
116 PR_POWER, /* '**' */
117 PR_AWAIT, /* 'await' */
118 PR_ATOM,
119};
120
121static int
122append_ast_boolop(_PyUnicodeWriter *writer, expr_ty e, int level)
123{
124 Py_ssize_t i, value_count;
125 asdl_expr_seq *values;
126 const char *op = (e->v.BoolOp.op == And) ? " and " : " or ";
127 int pr = (e->v.BoolOp.op == And) ? PR_AND : PR_OR;
128
129 APPEND_STR_IF(level > pr, "(");
130
131 values = e->v.BoolOp.values;
132 value_count = asdl_seq_LEN(values);
133
134 for (i = 0; i < value_count; ++i) {
135 APPEND_STR_IF(i > 0, op);
136 APPEND_EXPR((expr_ty)asdl_seq_GET(values, i), pr + 1);
137 }
138
139 APPEND_STR_IF(level > pr, ")");
140 return 0;
141}
142
143static int
144append_ast_binop(_PyUnicodeWriter *writer, expr_ty e, int level)
145{
146 const char *op;
147 int pr;
148 bool rassoc = false; /* is right-associative? */
149
150 switch (e->v.BinOp.op) {
151 case Add: op = " + "; pr = PR_ARITH; break;
152 case Sub: op = " - "; pr = PR_ARITH; break;
153 case Mult: op = " * "; pr = PR_TERM; break;
154 case MatMult: op = " @ "; pr = PR_TERM; break;
155 case Div: op = " / "; pr = PR_TERM; break;
156 case Mod: op = " % "; pr = PR_TERM; break;
157 case LShift: op = " << "; pr = PR_SHIFT; break;
158 case RShift: op = " >> "; pr = PR_SHIFT; break;
159 case BitOr: op = " | "; pr = PR_BOR; break;
160 case BitXor: op = " ^ "; pr = PR_BXOR; break;
161 case BitAnd: op = " & "; pr = PR_BAND; break;
162 case FloorDiv: op = " // "; pr = PR_TERM; break;
163 case Pow: op = " ** "; pr = PR_POWER; rassoc = true; break;
164 default:
165 PyErr_SetString(PyExc_SystemError,
166 "unknown binary operator");
167 return -1;
168 }
169
170 APPEND_STR_IF(level > pr, "(");
171 APPEND_EXPR(e->v.BinOp.left, pr + rassoc);
172 APPEND_STR(op);
173 APPEND_EXPR(e->v.BinOp.right, pr + !rassoc);
174 APPEND_STR_IF(level > pr, ")");
175 return 0;
176}
177
178static int
179append_ast_unaryop(_PyUnicodeWriter *writer, expr_ty e, int level)
180{
181 const char *op;
182 int pr;
183
184 switch (e->v.UnaryOp.op) {
185 case Invert: op = "~"; pr = PR_FACTOR; break;
186 case Not: op = "not "; pr = PR_NOT; break;
187 case UAdd: op = "+"; pr = PR_FACTOR; break;
188 case USub: op = "-"; pr = PR_FACTOR; break;
189 default:
190 PyErr_SetString(PyExc_SystemError,
191 "unknown unary operator");
192 return -1;
193 }
194
195 APPEND_STR_IF(level > pr, "(");
196 APPEND_STR(op);
197 APPEND_EXPR(e->v.UnaryOp.operand, pr);
198 APPEND_STR_IF(level > pr, ")");
199 return 0;
200}
201
202static int
203append_ast_arg(_PyUnicodeWriter *writer, arg_ty arg)
204{
205 if (-1 == _PyUnicodeWriter_WriteStr(writer, arg->arg)) {
206 return -1;
207 }
208 if (arg->annotation) {
209 APPEND_STR(": ");
210 APPEND_EXPR(arg->annotation, PR_TEST);
211 }
212 return 0;
213}
214
215static int
216append_ast_args(_PyUnicodeWriter *writer, arguments_ty args)
217{
218 bool first;
219 Py_ssize_t i, di, arg_count, posonlyarg_count, default_count;
220
221 first = true;
222
223 /* positional-only and positional arguments with defaults */
224 posonlyarg_count = asdl_seq_LEN(args->posonlyargs);
225 arg_count = asdl_seq_LEN(args->args);
226 default_count = asdl_seq_LEN(args->defaults);
227 for (i = 0; i < posonlyarg_count + arg_count; i++) {
228 APPEND_STR_IF_NOT_FIRST(", ");
229 if (i < posonlyarg_count){
230 APPEND(arg, (arg_ty)asdl_seq_GET(args->posonlyargs, i));
231 } else {
232 APPEND(arg, (arg_ty)asdl_seq_GET(args->args, i-posonlyarg_count));
233 }
234
235 di = i - posonlyarg_count - arg_count + default_count;
236 if (di >= 0) {
237 APPEND_STR("=");
238 APPEND_EXPR((expr_ty)asdl_seq_GET(args->defaults, di), PR_TEST);
239 }
240 if (posonlyarg_count && i + 1 == posonlyarg_count) {
241 APPEND_STR(", /");
242 }
243 }
244
245 /* vararg, or bare '*' if no varargs but keyword-only arguments present */
246 if (args->vararg || asdl_seq_LEN(args->kwonlyargs)) {
247 APPEND_STR_IF_NOT_FIRST(", ");
248 APPEND_STR("*");
249 if (args->vararg) {
250 APPEND(arg, args->vararg);
251 }
252 }
253
254 /* keyword-only arguments */
255 arg_count = asdl_seq_LEN(args->kwonlyargs);
256 default_count = asdl_seq_LEN(args->kw_defaults);
257 for (i = 0; i < arg_count; i++) {
258 APPEND_STR_IF_NOT_FIRST(", ");
259 APPEND(arg, (arg_ty)asdl_seq_GET(args->kwonlyargs, i));
260
261 di = i - arg_count + default_count;
262 if (di >= 0) {
263 expr_ty default_ = (expr_ty)asdl_seq_GET(args->kw_defaults, di);
264 if (default_) {
265 APPEND_STR("=");
266 APPEND_EXPR(default_, PR_TEST);
267 }
268 }
269 }
270
271 /* **kwargs */
272 if (args->kwarg) {
273 APPEND_STR_IF_NOT_FIRST(", ");
274 APPEND_STR("**");
275 APPEND(arg, args->kwarg);
276 }
277
278 return 0;
279}
280
281static int
282append_ast_lambda(_PyUnicodeWriter *writer, expr_ty e, int level)
283{
284 APPEND_STR_IF(level > PR_TEST, "(");
285 Py_ssize_t n_positional = (asdl_seq_LEN(e->v.Lambda.args->args) +
286 asdl_seq_LEN(e->v.Lambda.args->posonlyargs));
287 APPEND_STR(n_positional ? "lambda " : "lambda");
288 APPEND(args, e->v.Lambda.args);
289 APPEND_STR(": ");
290 APPEND_EXPR(e->v.Lambda.body, PR_TEST);
291 APPEND_STR_IF(level > PR_TEST, ")");
292 return 0;
293}
294
295static int
296append_ast_ifexp(_PyUnicodeWriter *writer, expr_ty e, int level)
297{
298 APPEND_STR_IF(level > PR_TEST, "(");
299 APPEND_EXPR(e->v.IfExp.body, PR_TEST + 1);
300 APPEND_STR(" if ");
301 APPEND_EXPR(e->v.IfExp.test, PR_TEST + 1);
302 APPEND_STR(" else ");
303 APPEND_EXPR(e->v.IfExp.orelse, PR_TEST);
304 APPEND_STR_IF(level > PR_TEST, ")");
305 return 0;
306}
307
308static int
309append_ast_dict(_PyUnicodeWriter *writer, expr_ty e)
310{
311 Py_ssize_t i, value_count;
312 expr_ty key_node;
313
314 APPEND_STR("{");
315 value_count = asdl_seq_LEN(e->v.Dict.values);
316
317 for (i = 0; i < value_count; i++) {
318 APPEND_STR_IF(i > 0, ", ");
319 key_node = (expr_ty)asdl_seq_GET(e->v.Dict.keys, i);
320 if (key_node != NULL) {
321 APPEND_EXPR(key_node, PR_TEST);
322 APPEND_STR(": ");
323 APPEND_EXPR((expr_ty)asdl_seq_GET(e->v.Dict.values, i), PR_TEST);
324 }
325 else {
326 APPEND_STR("**");
327 APPEND_EXPR((expr_ty)asdl_seq_GET(e->v.Dict.values, i), PR_EXPR);
328 }
329 }
330
331 APPEND_STR_FINISH("}");
332}
333
334static int
335append_ast_set(_PyUnicodeWriter *writer, expr_ty e)
336{
337 Py_ssize_t i, elem_count;
338
339 APPEND_STR("{");
340 elem_count = asdl_seq_LEN(e->v.Set.elts);
341 for (i = 0; i < elem_count; i++) {
342 APPEND_STR_IF(i > 0, ", ");
343 APPEND_EXPR((expr_ty)asdl_seq_GET(e->v.Set.elts, i), PR_TEST);
344 }
345
346 APPEND_STR_FINISH("}");
347}
348
349static int
350append_ast_list(_PyUnicodeWriter *writer, expr_ty e)
351{
352 Py_ssize_t i, elem_count;
353
354 APPEND_STR("[");
355 elem_count = asdl_seq_LEN(e->v.List.elts);
356 for (i = 0; i < elem_count; i++) {
357 APPEND_STR_IF(i > 0, ", ");
358 APPEND_EXPR((expr_ty)asdl_seq_GET(e->v.List.elts, i), PR_TEST);
359 }
360
361 APPEND_STR_FINISH("]");
362}
363
364static int
365append_ast_tuple(_PyUnicodeWriter *writer, expr_ty e, int level)
366{
367 Py_ssize_t i, elem_count;
368
369 elem_count = asdl_seq_LEN(e->v.Tuple.elts);
370
371 if (elem_count == 0) {
372 APPEND_STR_FINISH("()");
373 }
374
375 APPEND_STR_IF(level > PR_TUPLE, "(");
376
377 for (i = 0; i < elem_count; i++) {
378 APPEND_STR_IF(i > 0, ", ");
379 APPEND_EXPR((expr_ty)asdl_seq_GET(e->v.Tuple.elts, i), PR_TEST);
380 }
381
382 APPEND_STR_IF(elem_count == 1, ",");
383 APPEND_STR_IF(level > PR_TUPLE, ")");
384 return 0;
385}
386
387static int
388append_ast_comprehension(_PyUnicodeWriter *writer, comprehension_ty gen)
389{
390 Py_ssize_t i, if_count;
391
392 APPEND_STR(gen->is_async ? " async for " : " for ");
393 APPEND_EXPR(gen->target, PR_TUPLE);
394 APPEND_STR(" in ");
395 APPEND_EXPR(gen->iter, PR_TEST + 1);
396
397 if_count = asdl_seq_LEN(gen->ifs);
398 for (i = 0; i < if_count; i++) {
399 APPEND_STR(" if ");
400 APPEND_EXPR((expr_ty)asdl_seq_GET(gen->ifs, i), PR_TEST + 1);
401 }
402 return 0;
403}
404
405static int
406append_ast_comprehensions(_PyUnicodeWriter *writer, asdl_comprehension_seq *comprehensions)
407{
408 Py_ssize_t i, gen_count;
409 gen_count = asdl_seq_LEN(comprehensions);
410
411 for (i = 0; i < gen_count; i++) {
412 APPEND(comprehension, (comprehension_ty)asdl_seq_GET(comprehensions, i));
413 }
414
415 return 0;
416}
417
418static int
419append_ast_genexp(_PyUnicodeWriter *writer, expr_ty e)
420{
421 APPEND_STR("(");
422 APPEND_EXPR(e->v.GeneratorExp.elt, PR_TEST);
423 APPEND(comprehensions, e->v.GeneratorExp.generators);
424 APPEND_STR_FINISH(")");
425}
426
427static int
428append_ast_listcomp(_PyUnicodeWriter *writer, expr_ty e)
429{
430 APPEND_STR("[");
431 APPEND_EXPR(e->v.ListComp.elt, PR_TEST);
432 APPEND(comprehensions, e->v.ListComp.generators);
433 APPEND_STR_FINISH("]");
434}
435
436static int
437append_ast_setcomp(_PyUnicodeWriter *writer, expr_ty e)
438{
439 APPEND_STR("{");
440 APPEND_EXPR(e->v.SetComp.elt, PR_TEST);
441 APPEND(comprehensions, e->v.SetComp.generators);
442 APPEND_STR_FINISH("}");
443}
444
445static int
446append_ast_dictcomp(_PyUnicodeWriter *writer, expr_ty e)
447{
448 APPEND_STR("{");
449 APPEND_EXPR(e->v.DictComp.key, PR_TEST);
450 APPEND_STR(": ");
451 APPEND_EXPR(e->v.DictComp.value, PR_TEST);
452 APPEND(comprehensions, e->v.DictComp.generators);
453 APPEND_STR_FINISH("}");
454}
455
456static int
457append_ast_compare(_PyUnicodeWriter *writer, expr_ty e, int level)
458{
459 const char *op;
460 Py_ssize_t i, comparator_count;
461 asdl_expr_seq *comparators;
462 asdl_int_seq *ops;
463
464 APPEND_STR_IF(level > PR_CMP, "(");
465
466 comparators = e->v.Compare.comparators;
467 ops = e->v.Compare.ops;
468 comparator_count = asdl_seq_LEN(comparators);
469 assert(comparator_count > 0);
470 assert(comparator_count == asdl_seq_LEN(ops));
471
472 APPEND_EXPR(e->v.Compare.left, PR_CMP + 1);
473
474 for (i = 0; i < comparator_count; i++) {
475 switch ((cmpop_ty)asdl_seq_GET(ops, i)) {
476 case Eq:
477 op = " == ";
478 break;
479 case NotEq:
480 op = " != ";
481 break;
482 case Lt:
483 op = " < ";
484 break;
485 case LtE:
486 op = " <= ";
487 break;
488 case Gt:
489 op = " > ";
490 break;
491 case GtE:
492 op = " >= ";
493 break;
494 case Is:
495 op = " is ";
496 break;
497 case IsNot:
498 op = " is not ";
499 break;
500 case In:
501 op = " in ";
502 break;
503 case NotIn:
504 op = " not in ";
505 break;
506 default:
507 PyErr_SetString(PyExc_SystemError,
508 "unexpected comparison kind");
509 return -1;
510 }
511
512 APPEND_STR(op);
513 APPEND_EXPR((expr_ty)asdl_seq_GET(comparators, i), PR_CMP + 1);
514 }
515
516 APPEND_STR_IF(level > PR_CMP, ")");
517 return 0;
518}
519
520static int
521append_ast_keyword(_PyUnicodeWriter *writer, keyword_ty kw)
522{
523 if (kw->arg == NULL) {
524 APPEND_STR("**");
525 }
526 else {
527 if (-1 == _PyUnicodeWriter_WriteStr(writer, kw->arg)) {
528 return -1;
529 }
530
531 APPEND_STR("=");
532 }
533
534 APPEND_EXPR(kw->value, PR_TEST);
535 return 0;
536}
537
538static int
539append_ast_call(_PyUnicodeWriter *writer, expr_ty e)
540{
541 bool first;
542 Py_ssize_t i, arg_count, kw_count;
543 expr_ty expr;
544
545 APPEND_EXPR(e->v.Call.func, PR_ATOM);
546
547 arg_count = asdl_seq_LEN(e->v.Call.args);
548 kw_count = asdl_seq_LEN(e->v.Call.keywords);
549 if (arg_count == 1 && kw_count == 0) {
550 expr = (expr_ty)asdl_seq_GET(e->v.Call.args, 0);
551 if (expr->kind == GeneratorExp_kind) {
552 /* Special case: a single generator expression. */
553 return append_ast_genexp(writer, expr);
554 }
555 }
556
557 APPEND_STR("(");
558
559 first = true;
560 for (i = 0; i < arg_count; i++) {
561 APPEND_STR_IF_NOT_FIRST(", ");
562 APPEND_EXPR((expr_ty)asdl_seq_GET(e->v.Call.args, i), PR_TEST);
563 }
564
565 for (i = 0; i < kw_count; i++) {
566 APPEND_STR_IF_NOT_FIRST(", ");
567 APPEND(keyword, (keyword_ty)asdl_seq_GET(e->v.Call.keywords, i));
568 }
569
570 APPEND_STR_FINISH(")");
571}
572
573static PyObject *
574escape_braces(PyObject *orig)
575{
576 PyObject *temp;
577 PyObject *result;
578 temp = PyUnicode_Replace(orig, _str_open_br, _str_dbl_open_br, -1);
579 if (!temp) {
580 return NULL;
581 }
582 result = PyUnicode_Replace(temp, _str_close_br, _str_dbl_close_br, -1);
583 Py_DECREF(temp);
584 return result;
585}
586
587static int
588append_fstring_unicode(_PyUnicodeWriter *writer, PyObject *unicode)
589{
590 PyObject *escaped;
591 int result = -1;
592 escaped = escape_braces(unicode);
593 if (escaped) {
594 result = _PyUnicodeWriter_WriteStr(writer, escaped);
595 Py_DECREF(escaped);
596 }
597 return result;
598}
599
600static int
601append_fstring_element(_PyUnicodeWriter *writer, expr_ty e, bool is_format_spec)
602{
603 switch (e->kind) {
604 case Constant_kind:
605 return append_fstring_unicode(writer, e->v.Constant.value);
606 case JoinedStr_kind:
607 return append_joinedstr(writer, e, is_format_spec);
608 case FormattedValue_kind:
609 return append_formattedvalue(writer, e);
610 default:
611 PyErr_SetString(PyExc_SystemError,
612 "unknown expression kind inside f-string");
613 return -1;
614 }
615}
616
617/* Build body separately to enable wrapping the entire stream of Strs,
618 Constants and FormattedValues in one opening and one closing quote. */
619static PyObject *
620build_fstring_body(asdl_expr_seq *values, bool is_format_spec)
621{
622 Py_ssize_t i, value_count;
623 _PyUnicodeWriter body_writer;
624 _PyUnicodeWriter_Init(&body_writer);
625 body_writer.min_length = 256;
626 body_writer.overallocate = 1;
627
628 value_count = asdl_seq_LEN(values);
629 for (i = 0; i < value_count; ++i) {
630 if (-1 == append_fstring_element(&body_writer,
631 (expr_ty)asdl_seq_GET(values, i),
632 is_format_spec
633 )) {
634 _PyUnicodeWriter_Dealloc(&body_writer);
635 return NULL;
636 }
637 }
638
639 return _PyUnicodeWriter_Finish(&body_writer);
640}
641
642static int
643append_joinedstr(_PyUnicodeWriter *writer, expr_ty e, bool is_format_spec)
644{
645 int result = -1;
646 PyObject *body = build_fstring_body(e->v.JoinedStr.values, is_format_spec);
647 if (!body) {
648 return -1;
649 }
650
651 if (!is_format_spec) {
652 if (-1 != append_charp(writer, "f") &&
653 -1 != append_repr(writer, body))
654 {
655 result = 0;
656 }
657 }
658 else {
659 result = _PyUnicodeWriter_WriteStr(writer, body);
660 }
661 Py_DECREF(body);
662 return result;
663}
664
665static int
666append_formattedvalue(_PyUnicodeWriter *writer, expr_ty e)
667{
668 const char *conversion;
669 const char *outer_brace = "{";
670 /* Grammar allows PR_TUPLE, but use >PR_TEST for adding parenthesis
671 around a lambda with ':' */
672 PyObject *temp_fv_str = expr_as_unicode(e->v.FormattedValue.value, PR_TEST + 1);
673 if (!temp_fv_str) {
674 return -1;
675 }
676 if (PyUnicode_Find(temp_fv_str, _str_open_br, 0, 1, 1) == 0) {
677 /* Expression starts with a brace, split it with a space from the outer
678 one. */
679 outer_brace = "{ ";
680 }
681 if (-1 == append_charp(writer, outer_brace)) {
682 Py_DECREF(temp_fv_str);
683 return -1;
684 }
685 if (-1 == _PyUnicodeWriter_WriteStr(writer, temp_fv_str)) {
686 Py_DECREF(temp_fv_str);
687 return -1;
688 }
689 Py_DECREF(temp_fv_str);
690
691 if (e->v.FormattedValue.conversion > 0) {
692 switch (e->v.FormattedValue.conversion) {
693 case 'a':
694 conversion = "!a";
695 break;
696 case 'r':
697 conversion = "!r";
698 break;
699 case 's':
700 conversion = "!s";
701 break;
702 default:
703 PyErr_SetString(PyExc_SystemError,
704 "unknown f-value conversion kind");
705 return -1;
706 }
707 APPEND_STR(conversion);
708 }
709 if (e->v.FormattedValue.format_spec) {
710 if (-1 == _PyUnicodeWriter_WriteASCIIString(writer, ":", 1) ||
711 -1 == append_fstring_element(writer,
712 e->v.FormattedValue.format_spec,
713 true
714 ))
715 {
716 return -1;
717 }
718 }
719
720 APPEND_STR_FINISH("}");
721}
722
723static int
724append_ast_constant(_PyUnicodeWriter *writer, PyObject *constant)
725{
726 if (PyTuple_CheckExact(constant)) {
727 Py_ssize_t i, elem_count;
728
729 elem_count = PyTuple_GET_SIZE(constant);
730 APPEND_STR("(");
731 for (i = 0; i < elem_count; i++) {
732 APPEND_STR_IF(i > 0, ", ");
733 if (append_ast_constant(writer, PyTuple_GET_ITEM(constant, i)) < 0) {
734 return -1;
735 }
736 }
737
738 APPEND_STR_IF(elem_count == 1, ",");
739 APPEND_STR(")");
740 return 0;
741 }
742 return append_repr(writer, constant);
743}
744
745static int
746append_ast_attribute(_PyUnicodeWriter *writer, expr_ty e)
747{
748 const char *period;
749 expr_ty v = e->v.Attribute.value;
750 APPEND_EXPR(v, PR_ATOM);
751
752 /* Special case: integers require a space for attribute access to be
753 unambiguous. */
754 if (v->kind == Constant_kind && PyLong_CheckExact(v->v.Constant.value)) {
755 period = " .";
756 }
757 else {
758 period = ".";
759 }
760 APPEND_STR(period);
761
762 return _PyUnicodeWriter_WriteStr(writer, e->v.Attribute.attr);
763}
764
765static int
766append_ast_slice(_PyUnicodeWriter *writer, expr_ty e)
767{
768 if (e->v.Slice.lower) {
769 APPEND_EXPR(e->v.Slice.lower, PR_TEST);
770 }
771
772 APPEND_STR(":");
773
774 if (e->v.Slice.upper) {
775 APPEND_EXPR(e->v.Slice.upper, PR_TEST);
776 }
777
778 if (e->v.Slice.step) {
779 APPEND_STR(":");
780 APPEND_EXPR(e->v.Slice.step, PR_TEST);
781 }
782 return 0;
783}
784
785static int
786append_ast_subscript(_PyUnicodeWriter *writer, expr_ty e)
787{
788 APPEND_EXPR(e->v.Subscript.value, PR_ATOM);
789 int level = PR_TUPLE;
790 expr_ty slice = e->v.Subscript.slice;
791 if (slice->kind == Tuple_kind) {
792 for (Py_ssize_t i = 0; i < asdl_seq_LEN(slice->v.Tuple.elts); i++) {
793 expr_ty element = asdl_seq_GET(slice->v.Tuple.elts, i);
794 if (element->kind == Starred_kind) {
795 ++level;
796 break;
797 }
798 }
799 }
800 APPEND_STR("[");
801 APPEND_EXPR(e->v.Subscript.slice, level);
802 APPEND_STR_FINISH("]");
803}
804
805static int
806append_ast_starred(_PyUnicodeWriter *writer, expr_ty e)
807{
808 APPEND_STR("*");
809 APPEND_EXPR(e->v.Starred.value, PR_EXPR);
810 return 0;
811}
812
813static int
814append_ast_yield(_PyUnicodeWriter *writer, expr_ty e)
815{
816 if (!e->v.Yield.value) {
817 APPEND_STR_FINISH("(yield)");
818 }
819
820 APPEND_STR("(yield ");
821 APPEND_EXPR(e->v.Yield.value, PR_TEST);
822 APPEND_STR_FINISH(")");
823}
824
825static int
826append_ast_yield_from(_PyUnicodeWriter *writer, expr_ty e)
827{
828 APPEND_STR("(yield from ");
829 APPEND_EXPR(e->v.YieldFrom.value, PR_TEST);
830 APPEND_STR_FINISH(")");
831}
832
833static int
834append_ast_await(_PyUnicodeWriter *writer, expr_ty e, int level)
835{
836 APPEND_STR_IF(level > PR_AWAIT, "(");
837 APPEND_STR("await ");
838 APPEND_EXPR(e->v.Await.value, PR_ATOM);
839 APPEND_STR_IF(level > PR_AWAIT, ")");
840 return 0;
841}
842
843static int
844append_named_expr(_PyUnicodeWriter *writer, expr_ty e, int level)
845{
846 APPEND_STR_IF(level > PR_TUPLE, "(");
847 APPEND_EXPR(e->v.NamedExpr.target, PR_ATOM);
848 APPEND_STR(" := ");
849 APPEND_EXPR(e->v.NamedExpr.value, PR_ATOM);
850 APPEND_STR_IF(level > PR_TUPLE, ")");
851 return 0;
852}
853
854static int
855append_ast_expr(_PyUnicodeWriter *writer, expr_ty e, int level)
856{
857 switch (e->kind) {
858 case BoolOp_kind:
859 return append_ast_boolop(writer, e, level);
860 case BinOp_kind:
861 return append_ast_binop(writer, e, level);
862 case UnaryOp_kind:
863 return append_ast_unaryop(writer, e, level);
864 case Lambda_kind:
865 return append_ast_lambda(writer, e, level);
866 case IfExp_kind:
867 return append_ast_ifexp(writer, e, level);
868 case Dict_kind:
869 return append_ast_dict(writer, e);
870 case Set_kind:
871 return append_ast_set(writer, e);
872 case GeneratorExp_kind:
873 return append_ast_genexp(writer, e);
874 case ListComp_kind:
875 return append_ast_listcomp(writer, e);
876 case SetComp_kind:
877 return append_ast_setcomp(writer, e);
878 case DictComp_kind:
879 return append_ast_dictcomp(writer, e);
880 case Yield_kind:
881 return append_ast_yield(writer, e);
882 case YieldFrom_kind:
883 return append_ast_yield_from(writer, e);
884 case Await_kind:
885 return append_ast_await(writer, e, level);
886 case Compare_kind:
887 return append_ast_compare(writer, e, level);
888 case Call_kind:
889 return append_ast_call(writer, e);
890 case Constant_kind:
891 if (e->v.Constant.value == Py_Ellipsis) {
892 APPEND_STR_FINISH("...");
893 }
894 if (e->v.Constant.kind != NULL
895 && -1 == _PyUnicodeWriter_WriteStr(writer, e->v.Constant.kind)) {
896 return -1;
897 }
898 return append_ast_constant(writer, e->v.Constant.value);
899 case JoinedStr_kind:
900 return append_joinedstr(writer, e, false);
901 case FormattedValue_kind:
902 return append_formattedvalue(writer, e);
903 /* The following exprs can be assignment targets. */
904 case Attribute_kind:
905 return append_ast_attribute(writer, e);
906 case Subscript_kind:
907 return append_ast_subscript(writer, e);
908 case Starred_kind:
909 return append_ast_starred(writer, e);
910 case Slice_kind:
911 return append_ast_slice(writer, e);
912 case Name_kind:
913 return _PyUnicodeWriter_WriteStr(writer, e->v.Name.id);
914 case List_kind:
915 return append_ast_list(writer, e);
916 case Tuple_kind:
917 return append_ast_tuple(writer, e, level);
918 case NamedExpr_kind:
919 return append_named_expr(writer, e, level);
920 // No default so compiler emits a warning for unhandled cases
921 }
922 PyErr_SetString(PyExc_SystemError,
923 "unknown expression kind");
924 return -1;
925}
926
927static int
928maybe_init_static_strings(void)
929{
930 if (!_str_open_br &&
931 !(_str_open_br = PyUnicode_InternFromString("{"))) {
932 return -1;
933 }
934 if (!_str_dbl_open_br &&
935 !(_str_dbl_open_br = PyUnicode_InternFromString("{{"))) {
936 return -1;
937 }
938 if (!_str_close_br &&
939 !(_str_close_br = PyUnicode_InternFromString("}"))) {
940 return -1;
941 }
942 if (!_str_dbl_close_br &&
943 !(_str_dbl_close_br = PyUnicode_InternFromString("}}"))) {
944 return -1;
945 }
946 if (!_str_inf &&
947 !(_str_inf = PyUnicode_FromString("inf"))) {
948 return -1;
949 }
950 if (!_str_replace_inf &&
951 !(_str_replace_inf = PyUnicode_FromFormat("1e%d", 1 + DBL_MAX_10_EXP))) {
952 return -1;
953 }
954 return 0;
955}
956
957static PyObject *
958expr_as_unicode(expr_ty e, int level)
959{
960 _PyUnicodeWriter writer;
961 _PyUnicodeWriter_Init(&writer);
962 writer.min_length = 256;
963 writer.overallocate = 1;
964 if (-1 == maybe_init_static_strings() ||
965 -1 == append_ast_expr(&writer, e, level))
966 {
967 _PyUnicodeWriter_Dealloc(&writer);
968 return NULL;
969 }
970 return _PyUnicodeWriter_Finish(&writer);
971}
972
973PyObject *
974_PyAST_ExprAsUnicode(expr_ty e)
975{
976 return expr_as_unicode(e, PR_TEST);
977}
978