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 
11 static PyObject *_str_open_br;
12 static PyObject *_str_dbl_open_br;
13 static PyObject *_str_close_br;
14 static PyObject *_str_dbl_close_br;
15 static PyObject *_str_inf;
16 static PyObject *_str_replace_inf;
17 
18 /* Forward declarations for recursion via helper functions. */
19 static PyObject *
20 expr_as_unicode(expr_ty e, int level);
21 static int
22 append_ast_expr(_PyUnicodeWriter *writer, expr_ty e, int level);
23 static int
24 append_joinedstr(_PyUnicodeWriter *writer, expr_ty e, bool is_format_spec);
25 static int
26 append_formattedvalue(_PyUnicodeWriter *writer, expr_ty e);
27 static int
28 append_ast_slice(_PyUnicodeWriter *writer, expr_ty e);
29 
30 static int
append_charp(_PyUnicodeWriter * writer,const char * charp)31 append_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 
69 static int
append_repr(_PyUnicodeWriter * writer,PyObject * obj)70 append_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 
100 enum {
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 
121 static int
append_ast_boolop(_PyUnicodeWriter * writer,expr_ty e,int level)122 append_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 
143 static int
append_ast_binop(_PyUnicodeWriter * writer,expr_ty e,int level)144 append_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 
178 static int
append_ast_unaryop(_PyUnicodeWriter * writer,expr_ty e,int level)179 append_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 
202 static int
append_ast_arg(_PyUnicodeWriter * writer,arg_ty arg)203 append_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 
215 static int
append_ast_args(_PyUnicodeWriter * writer,arguments_ty args)216 append_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 
281 static int
append_ast_lambda(_PyUnicodeWriter * writer,expr_ty e,int level)282 append_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 
295 static int
append_ast_ifexp(_PyUnicodeWriter * writer,expr_ty e,int level)296 append_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 
308 static int
append_ast_dict(_PyUnicodeWriter * writer,expr_ty e)309 append_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 
334 static int
append_ast_set(_PyUnicodeWriter * writer,expr_ty e)335 append_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 
349 static int
append_ast_list(_PyUnicodeWriter * writer,expr_ty e)350 append_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 
364 static int
append_ast_tuple(_PyUnicodeWriter * writer,expr_ty e,int level)365 append_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 
387 static int
append_ast_comprehension(_PyUnicodeWriter * writer,comprehension_ty gen)388 append_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 
405 static int
append_ast_comprehensions(_PyUnicodeWriter * writer,asdl_comprehension_seq * comprehensions)406 append_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 
418 static int
append_ast_genexp(_PyUnicodeWriter * writer,expr_ty e)419 append_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 
427 static int
append_ast_listcomp(_PyUnicodeWriter * writer,expr_ty e)428 append_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 
436 static int
append_ast_setcomp(_PyUnicodeWriter * writer,expr_ty e)437 append_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 
445 static int
append_ast_dictcomp(_PyUnicodeWriter * writer,expr_ty e)446 append_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 
456 static int
append_ast_compare(_PyUnicodeWriter * writer,expr_ty e,int level)457 append_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 
520 static int
append_ast_keyword(_PyUnicodeWriter * writer,keyword_ty kw)521 append_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 
538 static int
append_ast_call(_PyUnicodeWriter * writer,expr_ty e)539 append_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 
573 static PyObject *
escape_braces(PyObject * orig)574 escape_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 
587 static int
append_fstring_unicode(_PyUnicodeWriter * writer,PyObject * unicode)588 append_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 
600 static int
append_fstring_element(_PyUnicodeWriter * writer,expr_ty e,bool is_format_spec)601 append_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. */
619 static PyObject *
build_fstring_body(asdl_expr_seq * values,bool is_format_spec)620 build_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 
642 static int
append_joinedstr(_PyUnicodeWriter * writer,expr_ty e,bool is_format_spec)643 append_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 
665 static int
append_formattedvalue(_PyUnicodeWriter * writer,expr_ty e)666 append_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 
723 static int
append_ast_constant(_PyUnicodeWriter * writer,PyObject * constant)724 append_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 
745 static int
append_ast_attribute(_PyUnicodeWriter * writer,expr_ty e)746 append_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 
765 static int
append_ast_slice(_PyUnicodeWriter * writer,expr_ty e)766 append_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 
785 static int
append_ast_subscript(_PyUnicodeWriter * writer,expr_ty e)786 append_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 
805 static int
append_ast_starred(_PyUnicodeWriter * writer,expr_ty e)806 append_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 
813 static int
append_ast_yield(_PyUnicodeWriter * writer,expr_ty e)814 append_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 
825 static int
append_ast_yield_from(_PyUnicodeWriter * writer,expr_ty e)826 append_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 
833 static int
append_ast_await(_PyUnicodeWriter * writer,expr_ty e,int level)834 append_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 
843 static int
append_named_expr(_PyUnicodeWriter * writer,expr_ty e,int level)844 append_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 
854 static int
append_ast_expr(_PyUnicodeWriter * writer,expr_ty e,int level)855 append_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 
927 static int
maybe_init_static_strings(void)928 maybe_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 
957 static PyObject *
expr_as_unicode(expr_ty e,int level)958 expr_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 
973 PyObject *
_PyAST_ExprAsUnicode(expr_ty e)974 _PyAST_ExprAsUnicode(expr_ty e)
975 {
976     return expr_as_unicode(e, PR_TEST);
977 }
978