• 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, default_count;
197 
198     first = true;
199 
200     /* positional arguments with defaults */
201     arg_count = asdl_seq_LEN(args->args);
202     default_count = asdl_seq_LEN(args->defaults);
203     for (i = 0; i < arg_count; i++) {
204         APPEND_STR_IF_NOT_FIRST(", ");
205         APPEND(arg, (arg_ty)asdl_seq_GET(args->args, i));
206 
207         di = i - arg_count + default_count;
208         if (di >= 0) {
209             APPEND_STR("=");
210             APPEND_EXPR((expr_ty)asdl_seq_GET(args->defaults, di), PR_TEST);
211         }
212     }
213 
214     /* vararg, or bare '*' if no varargs but keyword-only arguments present */
215     if (args->vararg || asdl_seq_LEN(args->kwonlyargs)) {
216         APPEND_STR_IF_NOT_FIRST(", ");
217         APPEND_STR("*");
218         if (args->vararg) {
219             APPEND(arg, args->vararg);
220         }
221     }
222 
223     /* keyword-only arguments */
224     arg_count = asdl_seq_LEN(args->kwonlyargs);
225     default_count = asdl_seq_LEN(args->kw_defaults);
226     for (i = 0; i < arg_count; i++) {
227         APPEND_STR_IF_NOT_FIRST(", ");
228         APPEND(arg, (arg_ty)asdl_seq_GET(args->kwonlyargs, i));
229 
230         di = i - arg_count + default_count;
231         if (di >= 0) {
232             expr_ty default_ = (expr_ty)asdl_seq_GET(args->kw_defaults, di);
233             if (default_) {
234                 APPEND_STR("=");
235                 APPEND_EXPR(default_, PR_TEST);
236             }
237         }
238     }
239 
240     /* **kwargs */
241     if (args->kwarg) {
242         APPEND_STR_IF_NOT_FIRST(", ");
243         APPEND_STR("**");
244         APPEND(arg, args->kwarg);
245     }
246 
247     return 0;
248 }
249 
250 static int
append_ast_lambda(_PyUnicodeWriter * writer,expr_ty e,int level)251 append_ast_lambda(_PyUnicodeWriter *writer, expr_ty e, int level)
252 {
253     APPEND_STR_IF(level > PR_TEST, "(");
254     APPEND_STR(asdl_seq_LEN(e->v.Lambda.args->args) ? "lambda " : "lambda");
255     APPEND(args, e->v.Lambda.args);
256     APPEND_STR(": ");
257     APPEND_EXPR(e->v.Lambda.body, PR_TEST);
258     APPEND_STR_IF(level > PR_TEST, ")");
259     return 0;
260 }
261 
262 static int
append_ast_ifexp(_PyUnicodeWriter * writer,expr_ty e,int level)263 append_ast_ifexp(_PyUnicodeWriter *writer, expr_ty e, int level)
264 {
265     APPEND_STR_IF(level > PR_TEST, "(");
266     APPEND_EXPR(e->v.IfExp.body, PR_TEST + 1);
267     APPEND_STR(" if ");
268     APPEND_EXPR(e->v.IfExp.test, PR_TEST + 1);
269     APPEND_STR(" else ");
270     APPEND_EXPR(e->v.IfExp.orelse, PR_TEST);
271     APPEND_STR_IF(level > PR_TEST, ")");
272     return 0;
273 }
274 
275 static int
append_ast_dict(_PyUnicodeWriter * writer,expr_ty e)276 append_ast_dict(_PyUnicodeWriter *writer, expr_ty e)
277 {
278     Py_ssize_t i, value_count;
279     expr_ty key_node;
280 
281     APPEND_STR("{");
282     value_count = asdl_seq_LEN(e->v.Dict.values);
283 
284     for (i = 0; i < value_count; i++) {
285         APPEND_STR_IF(i > 0, ", ");
286         key_node = (expr_ty)asdl_seq_GET(e->v.Dict.keys, i);
287         if (key_node != NULL) {
288             APPEND_EXPR(key_node, PR_TEST);
289             APPEND_STR(": ");
290             APPEND_EXPR((expr_ty)asdl_seq_GET(e->v.Dict.values, i), PR_TEST);
291         }
292         else {
293             APPEND_STR("**");
294             APPEND_EXPR((expr_ty)asdl_seq_GET(e->v.Dict.values, i), PR_EXPR);
295         }
296     }
297 
298     APPEND_STR_FINISH("}");
299 }
300 
301 static int
append_ast_set(_PyUnicodeWriter * writer,expr_ty e)302 append_ast_set(_PyUnicodeWriter *writer, expr_ty e)
303 {
304     Py_ssize_t i, elem_count;
305 
306     APPEND_STR("{");
307     elem_count = asdl_seq_LEN(e->v.Set.elts);
308     for (i = 0; i < elem_count; i++) {
309         APPEND_STR_IF(i > 0, ", ");
310         APPEND_EXPR((expr_ty)asdl_seq_GET(e->v.Set.elts, i), PR_TEST);
311     }
312 
313     APPEND_STR_FINISH("}");
314 }
315 
316 static int
append_ast_list(_PyUnicodeWriter * writer,expr_ty e)317 append_ast_list(_PyUnicodeWriter *writer, expr_ty e)
318 {
319     Py_ssize_t i, elem_count;
320 
321     APPEND_STR("[");
322     elem_count = asdl_seq_LEN(e->v.List.elts);
323     for (i = 0; i < elem_count; i++) {
324         APPEND_STR_IF(i > 0, ", ");
325         APPEND_EXPR((expr_ty)asdl_seq_GET(e->v.List.elts, i), PR_TEST);
326     }
327 
328     APPEND_STR_FINISH("]");
329 }
330 
331 static int
append_ast_tuple(_PyUnicodeWriter * writer,expr_ty e,int level)332 append_ast_tuple(_PyUnicodeWriter *writer, expr_ty e, int level)
333 {
334     Py_ssize_t i, elem_count;
335 
336     elem_count = asdl_seq_LEN(e->v.Tuple.elts);
337 
338     if (elem_count == 0) {
339         APPEND_STR_FINISH("()");
340     }
341 
342     APPEND_STR_IF(level > PR_TUPLE, "(");
343 
344     for (i = 0; i < elem_count; i++) {
345         APPEND_STR_IF(i > 0, ", ");
346         APPEND_EXPR((expr_ty)asdl_seq_GET(e->v.Tuple.elts, i), PR_TEST);
347     }
348 
349     APPEND_STR_IF(elem_count == 1, ",");
350     APPEND_STR_IF(level > PR_TUPLE, ")");
351     return 0;
352 }
353 
354 static int
append_ast_comprehension(_PyUnicodeWriter * writer,comprehension_ty gen)355 append_ast_comprehension(_PyUnicodeWriter *writer, comprehension_ty gen)
356 {
357     Py_ssize_t i, if_count;
358 
359     APPEND_STR(gen->is_async ? " async for " : " for ");
360     APPEND_EXPR(gen->target, PR_TUPLE);
361     APPEND_STR(" in ");
362     APPEND_EXPR(gen->iter, PR_TEST + 1);
363 
364     if_count = asdl_seq_LEN(gen->ifs);
365     for (i = 0; i < if_count; i++) {
366         APPEND_STR(" if ");
367         APPEND_EXPR((expr_ty)asdl_seq_GET(gen->ifs, i), PR_TEST + 1);
368     }
369     return 0;
370 }
371 
372 static int
append_ast_comprehensions(_PyUnicodeWriter * writer,asdl_seq * comprehensions)373 append_ast_comprehensions(_PyUnicodeWriter *writer, asdl_seq *comprehensions)
374 {
375     Py_ssize_t i, gen_count;
376     gen_count = asdl_seq_LEN(comprehensions);
377 
378     for (i = 0; i < gen_count; i++) {
379         APPEND(comprehension, (comprehension_ty)asdl_seq_GET(comprehensions, i));
380     }
381 
382     return 0;
383 }
384 
385 static int
append_ast_genexp(_PyUnicodeWriter * writer,expr_ty e)386 append_ast_genexp(_PyUnicodeWriter *writer, expr_ty e)
387 {
388     APPEND_STR("(");
389     APPEND_EXPR(e->v.GeneratorExp.elt, PR_TEST);
390     APPEND(comprehensions, e->v.GeneratorExp.generators);
391     APPEND_STR_FINISH(")");
392 }
393 
394 static int
append_ast_listcomp(_PyUnicodeWriter * writer,expr_ty e)395 append_ast_listcomp(_PyUnicodeWriter *writer, expr_ty e)
396 {
397     APPEND_STR("[");
398     APPEND_EXPR(e->v.ListComp.elt, PR_TEST);
399     APPEND(comprehensions, e->v.ListComp.generators);
400     APPEND_STR_FINISH("]");
401 }
402 
403 static int
append_ast_setcomp(_PyUnicodeWriter * writer,expr_ty e)404 append_ast_setcomp(_PyUnicodeWriter *writer, expr_ty e)
405 {
406     APPEND_STR("{");
407     APPEND_EXPR(e->v.SetComp.elt, PR_TEST);
408     APPEND(comprehensions, e->v.SetComp.generators);
409     APPEND_STR_FINISH("}");
410 }
411 
412 static int
append_ast_dictcomp(_PyUnicodeWriter * writer,expr_ty e)413 append_ast_dictcomp(_PyUnicodeWriter *writer, expr_ty e)
414 {
415     APPEND_STR("{");
416     APPEND_EXPR(e->v.DictComp.key, PR_TEST);
417     APPEND_STR(": ");
418     APPEND_EXPR(e->v.DictComp.value, PR_TEST);
419     APPEND(comprehensions, e->v.DictComp.generators);
420     APPEND_STR_FINISH("}");
421 }
422 
423 static int
append_ast_compare(_PyUnicodeWriter * writer,expr_ty e,int level)424 append_ast_compare(_PyUnicodeWriter *writer, expr_ty e, int level)
425 {
426     const char *op;
427     Py_ssize_t i, comparator_count;
428     asdl_seq *comparators;
429     asdl_int_seq *ops;
430 
431     APPEND_STR_IF(level > PR_CMP, "(");
432 
433     comparators = e->v.Compare.comparators;
434     ops = e->v.Compare.ops;
435     comparator_count = asdl_seq_LEN(comparators);
436     assert(comparator_count > 0);
437     assert(comparator_count == asdl_seq_LEN(ops));
438 
439     APPEND_EXPR(e->v.Compare.left, PR_CMP + 1);
440 
441     for (i = 0; i < comparator_count; i++) {
442         switch ((cmpop_ty)asdl_seq_GET(ops, i)) {
443         case Eq:
444             op = " == ";
445             break;
446         case NotEq:
447             op = " != ";
448             break;
449         case Lt:
450             op = " < ";
451             break;
452         case LtE:
453             op = " <= ";
454             break;
455         case Gt:
456             op = " > ";
457             break;
458         case GtE:
459             op = " >= ";
460             break;
461         case Is:
462             op = " is ";
463             break;
464         case IsNot:
465             op = " is not ";
466             break;
467         case In:
468             op = " in ";
469             break;
470         case NotIn:
471             op = " not in ";
472             break;
473         default:
474             PyErr_SetString(PyExc_SystemError,
475                             "unexpected comparison kind");
476             return -1;
477         }
478 
479         APPEND_STR(op);
480         APPEND_EXPR((expr_ty)asdl_seq_GET(comparators, i), PR_CMP + 1);
481     }
482 
483     APPEND_STR_IF(level > PR_CMP, ")");
484     return 0;
485 }
486 
487 static int
append_ast_keyword(_PyUnicodeWriter * writer,keyword_ty kw)488 append_ast_keyword(_PyUnicodeWriter *writer, keyword_ty kw)
489 {
490     if (kw->arg == NULL) {
491         APPEND_STR("**");
492     }
493     else {
494         if (-1 == _PyUnicodeWriter_WriteStr(writer, kw->arg)) {
495             return -1;
496         }
497 
498         APPEND_STR("=");
499     }
500 
501     APPEND_EXPR(kw->value, PR_TEST);
502     return 0;
503 }
504 
505 static int
append_ast_call(_PyUnicodeWriter * writer,expr_ty e)506 append_ast_call(_PyUnicodeWriter *writer, expr_ty e)
507 {
508     bool first;
509     Py_ssize_t i, arg_count, kw_count;
510     expr_ty expr;
511 
512     APPEND_EXPR(e->v.Call.func, PR_ATOM);
513 
514     arg_count = asdl_seq_LEN(e->v.Call.args);
515     kw_count = asdl_seq_LEN(e->v.Call.keywords);
516     if (arg_count == 1 && kw_count == 0) {
517         expr = (expr_ty)asdl_seq_GET(e->v.Call.args, 0);
518         if (expr->kind == GeneratorExp_kind) {
519             /* Special case: a single generator expression. */
520             return append_ast_genexp(writer, expr);
521         }
522     }
523 
524     APPEND_STR("(");
525 
526     first = true;
527     for (i = 0; i < arg_count; i++) {
528         APPEND_STR_IF_NOT_FIRST(", ");
529         APPEND_EXPR((expr_ty)asdl_seq_GET(e->v.Call.args, i), PR_TEST);
530     }
531 
532     for (i = 0; i < kw_count; i++) {
533         APPEND_STR_IF_NOT_FIRST(", ");
534         APPEND(keyword, (keyword_ty)asdl_seq_GET(e->v.Call.keywords, i));
535     }
536 
537     APPEND_STR_FINISH(")");
538 }
539 
540 static PyObject *
escape_braces(PyObject * orig)541 escape_braces(PyObject *orig)
542 {
543     PyObject *temp;
544     PyObject *result;
545     temp = PyUnicode_Replace(orig, _str_open_br, _str_dbl_open_br, -1);
546     if (!temp) {
547         return NULL;
548     }
549     result = PyUnicode_Replace(temp, _str_close_br, _str_dbl_close_br, -1);
550     Py_DECREF(temp);
551     return result;
552 }
553 
554 static int
append_fstring_unicode(_PyUnicodeWriter * writer,PyObject * unicode)555 append_fstring_unicode(_PyUnicodeWriter *writer, PyObject *unicode)
556 {
557     PyObject *escaped;
558     int result = -1;
559     escaped = escape_braces(unicode);
560     if (escaped) {
561         result = _PyUnicodeWriter_WriteStr(writer, escaped);
562         Py_DECREF(escaped);
563     }
564     return result;
565 }
566 
567 static int
append_fstring_element(_PyUnicodeWriter * writer,expr_ty e,bool is_format_spec)568 append_fstring_element(_PyUnicodeWriter *writer, expr_ty e, bool is_format_spec)
569 {
570     switch (e->kind) {
571     case Constant_kind:
572         return append_fstring_unicode(writer, e->v.Constant.value);
573     case Str_kind:
574         return append_fstring_unicode(writer, e->v.Str.s);
575     case JoinedStr_kind:
576         return append_joinedstr(writer, e, is_format_spec);
577     case FormattedValue_kind:
578         return append_formattedvalue(writer, e, is_format_spec);
579     default:
580         PyErr_SetString(PyExc_SystemError,
581                         "unknown expression kind inside f-string");
582         return -1;
583     }
584 }
585 
586 /* Build body separately to enable wrapping the entire stream of Strs,
587    Constants and FormattedValues in one opening and one closing quote. */
588 static PyObject *
build_fstring_body(asdl_seq * values,bool is_format_spec)589 build_fstring_body(asdl_seq *values, bool is_format_spec)
590 {
591     Py_ssize_t i, value_count;
592     _PyUnicodeWriter body_writer;
593     _PyUnicodeWriter_Init(&body_writer);
594     body_writer.min_length = 256;
595     body_writer.overallocate = 1;
596 
597     value_count = asdl_seq_LEN(values);
598     for (i = 0; i < value_count; ++i) {
599         if (-1 == append_fstring_element(&body_writer,
600                                          (expr_ty)asdl_seq_GET(values, i),
601                                          is_format_spec
602                                          )) {
603             _PyUnicodeWriter_Dealloc(&body_writer);
604             return NULL;
605         }
606     }
607 
608     return _PyUnicodeWriter_Finish(&body_writer);
609 }
610 
611 static int
append_joinedstr(_PyUnicodeWriter * writer,expr_ty e,bool is_format_spec)612 append_joinedstr(_PyUnicodeWriter *writer, expr_ty e, bool is_format_spec)
613 {
614     int result = -1;
615     PyObject *body = build_fstring_body(e->v.JoinedStr.values, is_format_spec);
616     if (!body) {
617         return -1;
618     }
619 
620     if (!is_format_spec) {
621         if (-1 != append_charp(writer, "f") &&
622             -1 != append_repr(writer, body))
623         {
624             result = 0;
625         }
626     }
627     else {
628         result = _PyUnicodeWriter_WriteStr(writer, body);
629     }
630     Py_DECREF(body);
631     return result;
632 }
633 
634 static int
append_formattedvalue(_PyUnicodeWriter * writer,expr_ty e,bool is_format_spec)635 append_formattedvalue(_PyUnicodeWriter *writer, expr_ty e, bool is_format_spec)
636 {
637     const char *conversion;
638     const char *outer_brace = "{";
639     /* Grammar allows PR_TUPLE, but use >PR_TEST for adding parenthesis
640        around a lambda with ':' */
641     PyObject *temp_fv_str = expr_as_unicode(e->v.FormattedValue.value, PR_TEST + 1);
642     if (!temp_fv_str) {
643         return -1;
644     }
645     if (PyUnicode_Find(temp_fv_str, _str_open_br, 0, 1, 1) == 0) {
646         /* Expression starts with a brace, split it with a space from the outer
647            one. */
648         outer_brace = "{ ";
649     }
650     if (-1 == append_charp(writer, outer_brace)) {
651         Py_DECREF(temp_fv_str);
652         return -1;
653     }
654     if (-1 == _PyUnicodeWriter_WriteStr(writer, temp_fv_str)) {
655         Py_DECREF(temp_fv_str);
656         return -1;
657     }
658     Py_DECREF(temp_fv_str);
659 
660     if (e->v.FormattedValue.conversion > 0) {
661         switch (e->v.FormattedValue.conversion) {
662         case 'a':
663             conversion = "!a";
664             break;
665         case 'r':
666             conversion = "!r";
667             break;
668         case 's':
669             conversion = "!s";
670             break;
671         default:
672             PyErr_SetString(PyExc_SystemError,
673                             "unknown f-value conversion kind");
674             return -1;
675         }
676         APPEND_STR(conversion);
677     }
678     if (e->v.FormattedValue.format_spec) {
679         if (-1 == _PyUnicodeWriter_WriteASCIIString(writer, ":", 1) ||
680             -1 == append_fstring_element(writer,
681                                          e->v.FormattedValue.format_spec,
682                                          true
683                                         ))
684         {
685             return -1;
686         }
687     }
688 
689     APPEND_STR_FINISH("}");
690 }
691 
692 static int
append_ast_attribute(_PyUnicodeWriter * writer,expr_ty e)693 append_ast_attribute(_PyUnicodeWriter *writer, expr_ty e)
694 {
695     const char *period;
696     APPEND_EXPR(e->v.Attribute.value, PR_ATOM);
697 
698     /* Special case: integers require a space for attribute access to be
699        unambiguous.  Floats and complex numbers don't but work with it, too. */
700     if (e->v.Attribute.value->kind == Num_kind ||
701         e->v.Attribute.value->kind == Constant_kind)
702     {
703         period = " .";
704     }
705     else {
706         period = ".";
707     }
708     APPEND_STR(period);
709 
710     return _PyUnicodeWriter_WriteStr(writer, e->v.Attribute.attr);
711 }
712 
713 static int
append_ast_simple_slice(_PyUnicodeWriter * writer,slice_ty slice)714 append_ast_simple_slice(_PyUnicodeWriter *writer, slice_ty slice)
715 {
716     if (slice->v.Slice.lower) {
717         APPEND_EXPR(slice->v.Slice.lower, PR_TEST);
718     }
719 
720     APPEND_STR(":");
721 
722     if (slice->v.Slice.upper) {
723         APPEND_EXPR(slice->v.Slice.upper, PR_TEST);
724     }
725 
726     if (slice->v.Slice.step) {
727         APPEND_STR(":");
728         APPEND_EXPR(slice->v.Slice.step, PR_TEST);
729     }
730     return 0;
731 }
732 
733 static int
append_ast_ext_slice(_PyUnicodeWriter * writer,slice_ty slice)734 append_ast_ext_slice(_PyUnicodeWriter *writer, slice_ty slice)
735 {
736     Py_ssize_t i, dims_count;
737     dims_count = asdl_seq_LEN(slice->v.ExtSlice.dims);
738     for (i = 0; i < dims_count; i++) {
739         APPEND_STR_IF(i > 0, ", ");
740         APPEND(slice, (slice_ty)asdl_seq_GET(slice->v.ExtSlice.dims, i));
741     }
742     return 0;
743 }
744 
745 static int
append_ast_slice(_PyUnicodeWriter * writer,slice_ty slice)746 append_ast_slice(_PyUnicodeWriter *writer, slice_ty slice)
747 {
748     switch (slice->kind) {
749     case Slice_kind:
750         return append_ast_simple_slice(writer, slice);
751     case ExtSlice_kind:
752         return append_ast_ext_slice(writer, slice);
753     case Index_kind:
754         APPEND_EXPR(slice->v.Index.value, PR_TUPLE);
755         return 0;
756     default:
757         PyErr_SetString(PyExc_SystemError,
758                         "unexpected slice kind");
759         return -1;
760     }
761 }
762 
763 static int
append_ast_subscript(_PyUnicodeWriter * writer,expr_ty e)764 append_ast_subscript(_PyUnicodeWriter *writer, expr_ty e)
765 {
766     APPEND_EXPR(e->v.Subscript.value, PR_ATOM);
767     APPEND_STR("[");
768     APPEND(slice, e->v.Subscript.slice);
769     APPEND_STR_FINISH("]");
770 }
771 
772 static int
append_ast_starred(_PyUnicodeWriter * writer,expr_ty e)773 append_ast_starred(_PyUnicodeWriter *writer, expr_ty e)
774 {
775     APPEND_STR("*");
776     APPEND_EXPR(e->v.Starred.value, PR_EXPR);
777     return 0;
778 }
779 
780 static int
append_ast_yield(_PyUnicodeWriter * writer,expr_ty e)781 append_ast_yield(_PyUnicodeWriter *writer, expr_ty e)
782 {
783     if (!e->v.Yield.value) {
784         APPEND_STR_FINISH("(yield)");
785     }
786 
787     APPEND_STR("(yield ");
788     APPEND_EXPR(e->v.Yield.value, PR_TEST);
789     APPEND_STR_FINISH(")");
790 }
791 
792 static int
append_ast_yield_from(_PyUnicodeWriter * writer,expr_ty e)793 append_ast_yield_from(_PyUnicodeWriter *writer, expr_ty e)
794 {
795     APPEND_STR("(yield from ");
796     APPEND_EXPR(e->v.YieldFrom.value, PR_TEST);
797     APPEND_STR_FINISH(")");
798 }
799 
800 static int
append_ast_await(_PyUnicodeWriter * writer,expr_ty e,int level)801 append_ast_await(_PyUnicodeWriter *writer, expr_ty e, int level)
802 {
803     APPEND_STR_IF(level > PR_AWAIT, "(");
804     APPEND_STR("await ");
805     APPEND_EXPR(e->v.Await.value, PR_ATOM);
806     APPEND_STR_IF(level > PR_AWAIT, ")");
807     return 0;
808 }
809 
810 static int
append_ast_expr(_PyUnicodeWriter * writer,expr_ty e,int level)811 append_ast_expr(_PyUnicodeWriter *writer, expr_ty e, int level)
812 {
813     switch (e->kind) {
814     case BoolOp_kind:
815         return append_ast_boolop(writer, e, level);
816     case BinOp_kind:
817         return append_ast_binop(writer, e, level);
818     case UnaryOp_kind:
819         return append_ast_unaryop(writer, e, level);
820     case Lambda_kind:
821         return append_ast_lambda(writer, e, level);
822     case IfExp_kind:
823         return append_ast_ifexp(writer, e, level);
824     case Dict_kind:
825         return append_ast_dict(writer, e);
826     case Set_kind:
827         return append_ast_set(writer, e);
828     case GeneratorExp_kind:
829         return append_ast_genexp(writer, e);
830     case ListComp_kind:
831         return append_ast_listcomp(writer, e);
832     case SetComp_kind:
833         return append_ast_setcomp(writer, e);
834     case DictComp_kind:
835         return append_ast_dictcomp(writer, e);
836     case Yield_kind:
837         return append_ast_yield(writer, e);
838     case YieldFrom_kind:
839         return append_ast_yield_from(writer, e);
840     case Await_kind:
841         return append_ast_await(writer, e, level);
842     case Compare_kind:
843         return append_ast_compare(writer, e, level);
844     case Call_kind:
845         return append_ast_call(writer, e);
846     case Constant_kind:
847         return append_repr(writer, e->v.Constant.value);
848     case Num_kind:
849         return append_repr(writer, e->v.Num.n);
850     case Str_kind:
851         return append_repr(writer, e->v.Str.s);
852     case JoinedStr_kind:
853         return append_joinedstr(writer, e, false);
854     case FormattedValue_kind:
855         return append_formattedvalue(writer, e, false);
856     case Bytes_kind:
857         return append_repr(writer, e->v.Bytes.s);
858     case Ellipsis_kind:
859         APPEND_STR_FINISH("...");
860     case NameConstant_kind:
861         return append_repr(writer, e->v.NameConstant.value);
862     /* The following exprs can be assignment targets. */
863     case Attribute_kind:
864         return append_ast_attribute(writer, e);
865     case Subscript_kind:
866         return append_ast_subscript(writer, e);
867     case Starred_kind:
868         return append_ast_starred(writer, e);
869     case Name_kind:
870         return _PyUnicodeWriter_WriteStr(writer, e->v.Name.id);
871     case List_kind:
872         return append_ast_list(writer, e);
873     case Tuple_kind:
874         return append_ast_tuple(writer, e, level);
875     default:
876         PyErr_SetString(PyExc_SystemError,
877                         "unknown expression kind");
878         return -1;
879     }
880 }
881 
882 static int
maybe_init_static_strings(void)883 maybe_init_static_strings(void)
884 {
885     if (!_str_open_br &&
886         !(_str_open_br = PyUnicode_InternFromString("{"))) {
887         return -1;
888     }
889     if (!_str_dbl_open_br &&
890         !(_str_dbl_open_br = PyUnicode_InternFromString("{{"))) {
891         return -1;
892     }
893     if (!_str_close_br &&
894         !(_str_close_br = PyUnicode_InternFromString("}"))) {
895         return -1;
896     }
897     if (!_str_dbl_close_br &&
898         !(_str_dbl_close_br = PyUnicode_InternFromString("}}"))) {
899         return -1;
900     }
901     return 0;
902 }
903 
904 static PyObject *
expr_as_unicode(expr_ty e,int level)905 expr_as_unicode(expr_ty e, int level)
906 {
907     _PyUnicodeWriter writer;
908     _PyUnicodeWriter_Init(&writer);
909     writer.min_length = 256;
910     writer.overallocate = 1;
911     if (-1 == maybe_init_static_strings() ||
912         -1 == append_ast_expr(&writer, e, level))
913     {
914         _PyUnicodeWriter_Dealloc(&writer);
915         return NULL;
916     }
917     return _PyUnicodeWriter_Finish(&writer);
918 }
919 
920 PyObject *
_PyAST_ExprAsUnicode(expr_ty e)921 _PyAST_ExprAsUnicode(expr_ty e)
922 {
923     return expr_as_unicode(e, PR_TEST);
924 }
925