• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 #include "Python.h"
2 #include "pycore_ast.h"           // expr_ty
3 #include "pycore_pystate.h"       // _PyInterpreterState_GET()
4 #include "pycore_runtime.h"       // _Py_ID()
5 #include <stdbool.h>
6 
7 /* This limited unparser is used to convert annotations back to strings
8  * during compilation rather than being a full AST unparser.
9  * See ast.unparse for a full unparser (written in Python)
10  */
11 
12 _Py_DECLARE_STR(dbl_open_br, "{{");
13 _Py_DECLARE_STR(dbl_close_br, "}}");
14 
15 /* Forward declarations for recursion via helper functions. */
16 static PyObject *
17 expr_as_unicode(expr_ty e, int level);
18 static int
19 append_ast_expr(_PyUnicodeWriter *writer, expr_ty e, int level);
20 static int
21 append_joinedstr(_PyUnicodeWriter *writer, expr_ty e, bool is_format_spec);
22 static int
23 append_formattedvalue(_PyUnicodeWriter *writer, expr_ty e);
24 static int
25 append_ast_slice(_PyUnicodeWriter *writer, expr_ty e);
26 
27 static int
append_charp(_PyUnicodeWriter * writer,const char * charp)28 append_charp(_PyUnicodeWriter *writer, const char *charp)
29 {
30     return _PyUnicodeWriter_WriteASCIIString(writer, charp, -1);
31 }
32 
33 #define APPEND_STR_FINISH(str)  do { \
34         return append_charp(writer, (str)); \
35     } while (0)
36 
37 #define APPEND_STR(str)  do { \
38         if (-1 == append_charp(writer, (str))) { \
39             return -1; \
40         } \
41     } while (0)
42 
43 #define APPEND_STR_IF(cond, str)  do { \
44         if ((cond) && -1 == append_charp(writer, (str))) { \
45             return -1; \
46         } \
47     } while (0)
48 
49 #define APPEND_STR_IF_NOT_FIRST(str)  do { \
50         APPEND_STR_IF(!first, (str)); \
51         first = false; \
52     } while (0)
53 
54 #define APPEND_EXPR(expr, pr)  do { \
55         if (-1 == append_ast_expr(writer, (expr), (pr))) { \
56             return -1; \
57         } \
58     } while (0)
59 
60 #define APPEND(type, value)  do { \
61         if (-1 == append_ast_ ## type(writer, (value))) { \
62             return -1; \
63         } \
64     } while (0)
65 
66 static int
append_repr(_PyUnicodeWriter * writer,PyObject * obj)67 append_repr(_PyUnicodeWriter *writer, PyObject *obj)
68 {
69     PyObject *repr = PyObject_Repr(obj);
70 
71     if (!repr) {
72         return -1;
73     }
74 
75     if ((PyFloat_CheckExact(obj) && Py_IS_INFINITY(PyFloat_AS_DOUBLE(obj))) ||
76         PyComplex_CheckExact(obj))
77     {
78         _Py_DECLARE_STR(str_replace_inf, "1e309");  // evaluates to inf
79         PyObject *new_repr = PyUnicode_Replace(
80             repr,
81             &_Py_ID(inf),
82             &_Py_STR(str_replace_inf),
83             -1
84         );
85         Py_DECREF(repr);
86         if (!new_repr) {
87             return -1;
88         }
89         repr = new_repr;
90     }
91     int ret = _PyUnicodeWriter_WriteStr(writer, repr);
92     Py_DECREF(repr);
93     return ret;
94 }
95 
96 /* Priority levels */
97 
98 enum {
99     PR_TUPLE,
100     PR_TEST,            /* 'if'-'else', 'lambda' */
101     PR_OR,              /* 'or' */
102     PR_AND,             /* 'and' */
103     PR_NOT,             /* 'not' */
104     PR_CMP,             /* '<', '>', '==', '>=', '<=', '!=',
105                            'in', 'not in', 'is', 'is not' */
106     PR_EXPR,
107     PR_BOR = PR_EXPR,   /* '|' */
108     PR_BXOR,            /* '^' */
109     PR_BAND,            /* '&' */
110     PR_SHIFT,           /* '<<', '>>' */
111     PR_ARITH,           /* '+', '-' */
112     PR_TERM,            /* '*', '@', '/', '%', '//' */
113     PR_FACTOR,          /* unary '+', '-', '~' */
114     PR_POWER,           /* '**' */
115     PR_AWAIT,           /* 'await' */
116     PR_ATOM,
117 };
118 
119 static int
append_ast_boolop(_PyUnicodeWriter * writer,expr_ty e,int level)120 append_ast_boolop(_PyUnicodeWriter *writer, expr_ty e, int level)
121 {
122     Py_ssize_t i, value_count;
123     asdl_expr_seq *values;
124     const char *op = (e->v.BoolOp.op == And) ? " and " : " or ";
125     int pr = (e->v.BoolOp.op == And) ? PR_AND : PR_OR;
126 
127     APPEND_STR_IF(level > pr, "(");
128 
129     values = e->v.BoolOp.values;
130     value_count = asdl_seq_LEN(values);
131 
132     for (i = 0; i < value_count; ++i) {
133         APPEND_STR_IF(i > 0, op);
134         APPEND_EXPR((expr_ty)asdl_seq_GET(values, i), pr + 1);
135     }
136 
137     APPEND_STR_IF(level > pr, ")");
138     return 0;
139 }
140 
141 static int
append_ast_binop(_PyUnicodeWriter * writer,expr_ty e,int level)142 append_ast_binop(_PyUnicodeWriter *writer, expr_ty e, int level)
143 {
144     const char *op;
145     int pr;
146     bool rassoc = false;  /* is right-associative? */
147 
148     switch (e->v.BinOp.op) {
149     case Add: op = " + "; pr = PR_ARITH; break;
150     case Sub: op = " - "; pr = PR_ARITH; break;
151     case Mult: op = " * "; pr = PR_TERM; break;
152     case MatMult: op = " @ "; pr = PR_TERM; break;
153     case Div: op = " / "; pr = PR_TERM; break;
154     case Mod: op = " % "; pr = PR_TERM; break;
155     case LShift: op = " << "; pr = PR_SHIFT; break;
156     case RShift: op = " >> "; pr = PR_SHIFT; break;
157     case BitOr: op = " | "; pr = PR_BOR; break;
158     case BitXor: op = " ^ "; pr = PR_BXOR; break;
159     case BitAnd: op = " & "; pr = PR_BAND; break;
160     case FloorDiv: op = " // "; pr = PR_TERM; break;
161     case Pow: op = " ** "; pr = PR_POWER; rassoc = true; break;
162     default:
163         PyErr_SetString(PyExc_SystemError,
164                         "unknown binary operator");
165         return -1;
166     }
167 
168     APPEND_STR_IF(level > pr, "(");
169     APPEND_EXPR(e->v.BinOp.left, pr + rassoc);
170     APPEND_STR(op);
171     APPEND_EXPR(e->v.BinOp.right, pr + !rassoc);
172     APPEND_STR_IF(level > pr, ")");
173     return 0;
174 }
175 
176 static int
append_ast_unaryop(_PyUnicodeWriter * writer,expr_ty e,int level)177 append_ast_unaryop(_PyUnicodeWriter *writer, expr_ty e, int level)
178 {
179     const char *op;
180     int pr;
181 
182     switch (e->v.UnaryOp.op) {
183     case Invert: op = "~"; pr = PR_FACTOR; break;
184     case Not: op = "not "; pr = PR_NOT; break;
185     case UAdd: op = "+"; pr = PR_FACTOR; break;
186     case USub: op = "-"; pr = PR_FACTOR; break;
187     default:
188         PyErr_SetString(PyExc_SystemError,
189                         "unknown unary operator");
190         return -1;
191     }
192 
193     APPEND_STR_IF(level > pr, "(");
194     APPEND_STR(op);
195     APPEND_EXPR(e->v.UnaryOp.operand, pr);
196     APPEND_STR_IF(level > pr, ")");
197     return 0;
198 }
199 
200 static int
append_ast_arg(_PyUnicodeWriter * writer,arg_ty arg)201 append_ast_arg(_PyUnicodeWriter *writer, arg_ty arg)
202 {
203     if (-1 == _PyUnicodeWriter_WriteStr(writer, arg->arg)) {
204         return -1;
205     }
206     if (arg->annotation) {
207         APPEND_STR(": ");
208         APPEND_EXPR(arg->annotation, PR_TEST);
209     }
210     return 0;
211 }
212 
213 static int
append_ast_args(_PyUnicodeWriter * writer,arguments_ty args)214 append_ast_args(_PyUnicodeWriter *writer, arguments_ty args)
215 {
216     bool first;
217     Py_ssize_t i, di, arg_count, posonlyarg_count, default_count;
218 
219     first = true;
220 
221     /* positional-only and positional arguments with defaults */
222     posonlyarg_count = asdl_seq_LEN(args->posonlyargs);
223     arg_count = asdl_seq_LEN(args->args);
224     default_count = asdl_seq_LEN(args->defaults);
225     for (i = 0; i < posonlyarg_count + arg_count; i++) {
226         APPEND_STR_IF_NOT_FIRST(", ");
227         if (i < posonlyarg_count){
228             APPEND(arg, (arg_ty)asdl_seq_GET(args->posonlyargs, i));
229         } else {
230             APPEND(arg, (arg_ty)asdl_seq_GET(args->args, i-posonlyarg_count));
231         }
232 
233         di = i - posonlyarg_count - arg_count + default_count;
234         if (di >= 0) {
235             APPEND_STR("=");
236             APPEND_EXPR((expr_ty)asdl_seq_GET(args->defaults, di), PR_TEST);
237         }
238         if (posonlyarg_count && i + 1 == posonlyarg_count) {
239             APPEND_STR(", /");
240         }
241     }
242 
243     /* vararg, or bare '*' if no varargs but keyword-only arguments present */
244     if (args->vararg || asdl_seq_LEN(args->kwonlyargs)) {
245         APPEND_STR_IF_NOT_FIRST(", ");
246         APPEND_STR("*");
247         if (args->vararg) {
248             APPEND(arg, args->vararg);
249         }
250     }
251 
252     /* keyword-only arguments */
253     arg_count = asdl_seq_LEN(args->kwonlyargs);
254     default_count = asdl_seq_LEN(args->kw_defaults);
255     for (i = 0; i < arg_count; i++) {
256         APPEND_STR_IF_NOT_FIRST(", ");
257         APPEND(arg, (arg_ty)asdl_seq_GET(args->kwonlyargs, i));
258 
259         di = i - arg_count + default_count;
260         if (di >= 0) {
261             expr_ty default_ = (expr_ty)asdl_seq_GET(args->kw_defaults, di);
262             if (default_) {
263                 APPEND_STR("=");
264                 APPEND_EXPR(default_, PR_TEST);
265             }
266         }
267     }
268 
269     /* **kwargs */
270     if (args->kwarg) {
271         APPEND_STR_IF_NOT_FIRST(", ");
272         APPEND_STR("**");
273         APPEND(arg, args->kwarg);
274     }
275 
276     return 0;
277 }
278 
279 static int
append_ast_lambda(_PyUnicodeWriter * writer,expr_ty e,int level)280 append_ast_lambda(_PyUnicodeWriter *writer, expr_ty e, int level)
281 {
282     APPEND_STR_IF(level > PR_TEST, "(");
283     Py_ssize_t n_positional = (asdl_seq_LEN(e->v.Lambda.args->args) +
284                                asdl_seq_LEN(e->v.Lambda.args->posonlyargs));
285     APPEND_STR(n_positional ? "lambda " : "lambda");
286     APPEND(args, e->v.Lambda.args);
287     APPEND_STR(": ");
288     APPEND_EXPR(e->v.Lambda.body, PR_TEST);
289     APPEND_STR_IF(level > PR_TEST, ")");
290     return 0;
291 }
292 
293 static int
append_ast_ifexp(_PyUnicodeWriter * writer,expr_ty e,int level)294 append_ast_ifexp(_PyUnicodeWriter *writer, expr_ty e, int level)
295 {
296     APPEND_STR_IF(level > PR_TEST, "(");
297     APPEND_EXPR(e->v.IfExp.body, PR_TEST + 1);
298     APPEND_STR(" if ");
299     APPEND_EXPR(e->v.IfExp.test, PR_TEST + 1);
300     APPEND_STR(" else ");
301     APPEND_EXPR(e->v.IfExp.orelse, PR_TEST);
302     APPEND_STR_IF(level > PR_TEST, ")");
303     return 0;
304 }
305 
306 static int
append_ast_dict(_PyUnicodeWriter * writer,expr_ty e)307 append_ast_dict(_PyUnicodeWriter *writer, expr_ty e)
308 {
309     Py_ssize_t i, value_count;
310     expr_ty key_node;
311 
312     APPEND_STR("{");
313     value_count = asdl_seq_LEN(e->v.Dict.values);
314 
315     for (i = 0; i < value_count; i++) {
316         APPEND_STR_IF(i > 0, ", ");
317         key_node = (expr_ty)asdl_seq_GET(e->v.Dict.keys, i);
318         if (key_node != NULL) {
319             APPEND_EXPR(key_node, PR_TEST);
320             APPEND_STR(": ");
321             APPEND_EXPR((expr_ty)asdl_seq_GET(e->v.Dict.values, i), PR_TEST);
322         }
323         else {
324             APPEND_STR("**");
325             APPEND_EXPR((expr_ty)asdl_seq_GET(e->v.Dict.values, i), PR_EXPR);
326         }
327     }
328 
329     APPEND_STR_FINISH("}");
330 }
331 
332 static int
append_ast_set(_PyUnicodeWriter * writer,expr_ty e)333 append_ast_set(_PyUnicodeWriter *writer, expr_ty e)
334 {
335     Py_ssize_t i, elem_count;
336 
337     APPEND_STR("{");
338     elem_count = asdl_seq_LEN(e->v.Set.elts);
339     for (i = 0; i < elem_count; i++) {
340         APPEND_STR_IF(i > 0, ", ");
341         APPEND_EXPR((expr_ty)asdl_seq_GET(e->v.Set.elts, i), PR_TEST);
342     }
343 
344     APPEND_STR_FINISH("}");
345 }
346 
347 static int
append_ast_list(_PyUnicodeWriter * writer,expr_ty e)348 append_ast_list(_PyUnicodeWriter *writer, expr_ty e)
349 {
350     Py_ssize_t i, elem_count;
351 
352     APPEND_STR("[");
353     elem_count = asdl_seq_LEN(e->v.List.elts);
354     for (i = 0; i < elem_count; i++) {
355         APPEND_STR_IF(i > 0, ", ");
356         APPEND_EXPR((expr_ty)asdl_seq_GET(e->v.List.elts, i), PR_TEST);
357     }
358 
359     APPEND_STR_FINISH("]");
360 }
361 
362 static int
append_ast_tuple(_PyUnicodeWriter * writer,expr_ty e,int level)363 append_ast_tuple(_PyUnicodeWriter *writer, expr_ty e, int level)
364 {
365     Py_ssize_t i, elem_count;
366 
367     elem_count = asdl_seq_LEN(e->v.Tuple.elts);
368 
369     if (elem_count == 0) {
370         APPEND_STR_FINISH("()");
371     }
372 
373     APPEND_STR_IF(level > PR_TUPLE, "(");
374 
375     for (i = 0; i < elem_count; i++) {
376         APPEND_STR_IF(i > 0, ", ");
377         APPEND_EXPR((expr_ty)asdl_seq_GET(e->v.Tuple.elts, i), PR_TEST);
378     }
379 
380     APPEND_STR_IF(elem_count == 1, ",");
381     APPEND_STR_IF(level > PR_TUPLE, ")");
382     return 0;
383 }
384 
385 static int
append_ast_comprehension(_PyUnicodeWriter * writer,comprehension_ty gen)386 append_ast_comprehension(_PyUnicodeWriter *writer, comprehension_ty gen)
387 {
388     Py_ssize_t i, if_count;
389 
390     APPEND_STR(gen->is_async ? " async for " : " for ");
391     APPEND_EXPR(gen->target, PR_TUPLE);
392     APPEND_STR(" in ");
393     APPEND_EXPR(gen->iter, PR_TEST + 1);
394 
395     if_count = asdl_seq_LEN(gen->ifs);
396     for (i = 0; i < if_count; i++) {
397         APPEND_STR(" if ");
398         APPEND_EXPR((expr_ty)asdl_seq_GET(gen->ifs, i), PR_TEST + 1);
399     }
400     return 0;
401 }
402 
403 static int
append_ast_comprehensions(_PyUnicodeWriter * writer,asdl_comprehension_seq * comprehensions)404 append_ast_comprehensions(_PyUnicodeWriter *writer, asdl_comprehension_seq *comprehensions)
405 {
406     Py_ssize_t i, gen_count;
407     gen_count = asdl_seq_LEN(comprehensions);
408 
409     for (i = 0; i < gen_count; i++) {
410         APPEND(comprehension, (comprehension_ty)asdl_seq_GET(comprehensions, i));
411     }
412 
413     return 0;
414 }
415 
416 static int
append_ast_genexp(_PyUnicodeWriter * writer,expr_ty e)417 append_ast_genexp(_PyUnicodeWriter *writer, expr_ty e)
418 {
419     APPEND_STR("(");
420     APPEND_EXPR(e->v.GeneratorExp.elt, PR_TEST);
421     APPEND(comprehensions, e->v.GeneratorExp.generators);
422     APPEND_STR_FINISH(")");
423 }
424 
425 static int
append_ast_listcomp(_PyUnicodeWriter * writer,expr_ty e)426 append_ast_listcomp(_PyUnicodeWriter *writer, expr_ty e)
427 {
428     APPEND_STR("[");
429     APPEND_EXPR(e->v.ListComp.elt, PR_TEST);
430     APPEND(comprehensions, e->v.ListComp.generators);
431     APPEND_STR_FINISH("]");
432 }
433 
434 static int
append_ast_setcomp(_PyUnicodeWriter * writer,expr_ty e)435 append_ast_setcomp(_PyUnicodeWriter *writer, expr_ty e)
436 {
437     APPEND_STR("{");
438     APPEND_EXPR(e->v.SetComp.elt, PR_TEST);
439     APPEND(comprehensions, e->v.SetComp.generators);
440     APPEND_STR_FINISH("}");
441 }
442 
443 static int
append_ast_dictcomp(_PyUnicodeWriter * writer,expr_ty e)444 append_ast_dictcomp(_PyUnicodeWriter *writer, expr_ty e)
445 {
446     APPEND_STR("{");
447     APPEND_EXPR(e->v.DictComp.key, PR_TEST);
448     APPEND_STR(": ");
449     APPEND_EXPR(e->v.DictComp.value, PR_TEST);
450     APPEND(comprehensions, e->v.DictComp.generators);
451     APPEND_STR_FINISH("}");
452 }
453 
454 static int
append_ast_compare(_PyUnicodeWriter * writer,expr_ty e,int level)455 append_ast_compare(_PyUnicodeWriter *writer, expr_ty e, int level)
456 {
457     const char *op;
458     Py_ssize_t i, comparator_count;
459     asdl_expr_seq *comparators;
460     asdl_int_seq *ops;
461 
462     APPEND_STR_IF(level > PR_CMP, "(");
463 
464     comparators = e->v.Compare.comparators;
465     ops = e->v.Compare.ops;
466     comparator_count = asdl_seq_LEN(comparators);
467     assert(comparator_count > 0);
468     assert(comparator_count == asdl_seq_LEN(ops));
469 
470     APPEND_EXPR(e->v.Compare.left, PR_CMP + 1);
471 
472     for (i = 0; i < comparator_count; i++) {
473         switch ((cmpop_ty)asdl_seq_GET(ops, i)) {
474         case Eq:
475             op = " == ";
476             break;
477         case NotEq:
478             op = " != ";
479             break;
480         case Lt:
481             op = " < ";
482             break;
483         case LtE:
484             op = " <= ";
485             break;
486         case Gt:
487             op = " > ";
488             break;
489         case GtE:
490             op = " >= ";
491             break;
492         case Is:
493             op = " is ";
494             break;
495         case IsNot:
496             op = " is not ";
497             break;
498         case In:
499             op = " in ";
500             break;
501         case NotIn:
502             op = " not in ";
503             break;
504         default:
505             PyErr_SetString(PyExc_SystemError,
506                             "unexpected comparison kind");
507             return -1;
508         }
509 
510         APPEND_STR(op);
511         APPEND_EXPR((expr_ty)asdl_seq_GET(comparators, i), PR_CMP + 1);
512     }
513 
514     APPEND_STR_IF(level > PR_CMP, ")");
515     return 0;
516 }
517 
518 static int
append_ast_keyword(_PyUnicodeWriter * writer,keyword_ty kw)519 append_ast_keyword(_PyUnicodeWriter *writer, keyword_ty kw)
520 {
521     if (kw->arg == NULL) {
522         APPEND_STR("**");
523     }
524     else {
525         if (-1 == _PyUnicodeWriter_WriteStr(writer, kw->arg)) {
526             return -1;
527         }
528 
529         APPEND_STR("=");
530     }
531 
532     APPEND_EXPR(kw->value, PR_TEST);
533     return 0;
534 }
535 
536 static int
append_ast_call(_PyUnicodeWriter * writer,expr_ty e)537 append_ast_call(_PyUnicodeWriter *writer, expr_ty e)
538 {
539     bool first;
540     Py_ssize_t i, arg_count, kw_count;
541     expr_ty expr;
542 
543     APPEND_EXPR(e->v.Call.func, PR_ATOM);
544 
545     arg_count = asdl_seq_LEN(e->v.Call.args);
546     kw_count = asdl_seq_LEN(e->v.Call.keywords);
547     if (arg_count == 1 && kw_count == 0) {
548         expr = (expr_ty)asdl_seq_GET(e->v.Call.args, 0);
549         if (expr->kind == GeneratorExp_kind) {
550             /* Special case: a single generator expression. */
551             return append_ast_genexp(writer, expr);
552         }
553     }
554 
555     APPEND_STR("(");
556 
557     first = true;
558     for (i = 0; i < arg_count; i++) {
559         APPEND_STR_IF_NOT_FIRST(", ");
560         APPEND_EXPR((expr_ty)asdl_seq_GET(e->v.Call.args, i), PR_TEST);
561     }
562 
563     for (i = 0; i < kw_count; i++) {
564         APPEND_STR_IF_NOT_FIRST(", ");
565         APPEND(keyword, (keyword_ty)asdl_seq_GET(e->v.Call.keywords, i));
566     }
567 
568     APPEND_STR_FINISH(")");
569 }
570 
571 static PyObject *
escape_braces(PyObject * orig)572 escape_braces(PyObject *orig)
573 {
574     PyObject *temp;
575     PyObject *result;
576     temp = PyUnicode_Replace(orig, _Py_LATIN1_CHR('{'),
577                              &_Py_STR(dbl_open_br), -1);
578     if (!temp) {
579         return NULL;
580     }
581     result = PyUnicode_Replace(temp, _Py_LATIN1_CHR('}'),
582                                &_Py_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, _Py_LATIN1_CHR('{'), 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     APPEND_STR("[");
790     APPEND_EXPR(e->v.Subscript.slice, PR_TUPLE);
791     APPEND_STR_FINISH("]");
792 }
793 
794 static int
append_ast_starred(_PyUnicodeWriter * writer,expr_ty e)795 append_ast_starred(_PyUnicodeWriter *writer, expr_ty e)
796 {
797     APPEND_STR("*");
798     APPEND_EXPR(e->v.Starred.value, PR_EXPR);
799     return 0;
800 }
801 
802 static int
append_ast_yield(_PyUnicodeWriter * writer,expr_ty e)803 append_ast_yield(_PyUnicodeWriter *writer, expr_ty e)
804 {
805     if (!e->v.Yield.value) {
806         APPEND_STR_FINISH("(yield)");
807     }
808 
809     APPEND_STR("(yield ");
810     APPEND_EXPR(e->v.Yield.value, PR_TEST);
811     APPEND_STR_FINISH(")");
812 }
813 
814 static int
append_ast_yield_from(_PyUnicodeWriter * writer,expr_ty e)815 append_ast_yield_from(_PyUnicodeWriter *writer, expr_ty e)
816 {
817     APPEND_STR("(yield from ");
818     APPEND_EXPR(e->v.YieldFrom.value, PR_TEST);
819     APPEND_STR_FINISH(")");
820 }
821 
822 static int
append_ast_await(_PyUnicodeWriter * writer,expr_ty e,int level)823 append_ast_await(_PyUnicodeWriter *writer, expr_ty e, int level)
824 {
825     APPEND_STR_IF(level > PR_AWAIT, "(");
826     APPEND_STR("await ");
827     APPEND_EXPR(e->v.Await.value, PR_ATOM);
828     APPEND_STR_IF(level > PR_AWAIT, ")");
829     return 0;
830 }
831 
832 static int
append_named_expr(_PyUnicodeWriter * writer,expr_ty e,int level)833 append_named_expr(_PyUnicodeWriter *writer, expr_ty e, int level)
834 {
835     APPEND_STR_IF(level > PR_TUPLE, "(");
836     APPEND_EXPR(e->v.NamedExpr.target, PR_ATOM);
837     APPEND_STR(" := ");
838     APPEND_EXPR(e->v.NamedExpr.value, PR_ATOM);
839     APPEND_STR_IF(level > PR_TUPLE, ")");
840     return 0;
841 }
842 
843 static int
append_ast_expr(_PyUnicodeWriter * writer,expr_ty e,int level)844 append_ast_expr(_PyUnicodeWriter *writer, expr_ty e, int level)
845 {
846     switch (e->kind) {
847     case BoolOp_kind:
848         return append_ast_boolop(writer, e, level);
849     case BinOp_kind:
850         return append_ast_binop(writer, e, level);
851     case UnaryOp_kind:
852         return append_ast_unaryop(writer, e, level);
853     case Lambda_kind:
854         return append_ast_lambda(writer, e, level);
855     case IfExp_kind:
856         return append_ast_ifexp(writer, e, level);
857     case Dict_kind:
858         return append_ast_dict(writer, e);
859     case Set_kind:
860         return append_ast_set(writer, e);
861     case GeneratorExp_kind:
862         return append_ast_genexp(writer, e);
863     case ListComp_kind:
864         return append_ast_listcomp(writer, e);
865     case SetComp_kind:
866         return append_ast_setcomp(writer, e);
867     case DictComp_kind:
868         return append_ast_dictcomp(writer, e);
869     case Yield_kind:
870         return append_ast_yield(writer, e);
871     case YieldFrom_kind:
872         return append_ast_yield_from(writer, e);
873     case Await_kind:
874         return append_ast_await(writer, e, level);
875     case Compare_kind:
876         return append_ast_compare(writer, e, level);
877     case Call_kind:
878         return append_ast_call(writer, e);
879     case Constant_kind:
880         if (e->v.Constant.value == Py_Ellipsis) {
881             APPEND_STR_FINISH("...");
882         }
883         if (e->v.Constant.kind != NULL
884             && -1 == _PyUnicodeWriter_WriteStr(writer, e->v.Constant.kind)) {
885             return -1;
886         }
887         return append_ast_constant(writer, e->v.Constant.value);
888     case JoinedStr_kind:
889         return append_joinedstr(writer, e, false);
890     case FormattedValue_kind:
891         return append_formattedvalue(writer, e);
892     /* The following exprs can be assignment targets. */
893     case Attribute_kind:
894         return append_ast_attribute(writer, e);
895     case Subscript_kind:
896         return append_ast_subscript(writer, e);
897     case Starred_kind:
898         return append_ast_starred(writer, e);
899     case Slice_kind:
900         return append_ast_slice(writer, e);
901     case Name_kind:
902         return _PyUnicodeWriter_WriteStr(writer, e->v.Name.id);
903     case List_kind:
904         return append_ast_list(writer, e);
905     case Tuple_kind:
906         return append_ast_tuple(writer, e, level);
907     case NamedExpr_kind:
908         return append_named_expr(writer, e, level);
909     // No default so compiler emits a warning for unhandled cases
910     }
911     PyErr_SetString(PyExc_SystemError,
912                     "unknown expression kind");
913     return -1;
914 }
915 
916 static PyObject *
expr_as_unicode(expr_ty e,int level)917 expr_as_unicode(expr_ty e, int level)
918 {
919     _PyUnicodeWriter writer;
920     _PyUnicodeWriter_Init(&writer);
921     writer.min_length = 256;
922     writer.overallocate = 1;
923     if (-1 == append_ast_expr(&writer, e, level)) {
924         _PyUnicodeWriter_Dealloc(&writer);
925         return NULL;
926     }
927     return _PyUnicodeWriter_Finish(&writer);
928 }
929 
930 PyObject *
_PyAST_ExprAsUnicode(expr_ty e)931 _PyAST_ExprAsUnicode(expr_ty e)
932 {
933     return expr_as_unicode(e, PR_TEST);
934 }
935