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