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