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