python_parser/visitors/
printer.rs

1//! Prints the AST as Python code.
2
3use super::super::ast::*;
4
5fn comma_join<'a, T2: ToString, T: IntoIterator<Item = T2>>(i: T) -> String {
6    let mut i = i.into_iter();
7    let mut s: String = i.next().map(|s| s.to_string()).unwrap_or("".to_string());
8    for s2 in i {
9        s.push_str(", ");
10        s.push_str(&s2.to_string()[..]);
11    }
12    s
13}
14
15fn space_join<'a, T2: ToString, T: IntoIterator<Item = T2>>(i: T) -> String {
16    let mut i = i.into_iter();
17    let mut s: String = i.next().map(|s| s.to_string()).unwrap_or("".to_string());
18    for s2 in i {
19        s.push_str(" ");
20        s.push_str(&s2.to_string()[..]);
21    }
22    s
23}
24
25fn dot_join<'a, T2: ToString, T: IntoIterator<Item = T2>>(i: T) -> String {
26    let mut i = i.into_iter();
27    let mut s: String = i.next().map(|s| s.to_string()).unwrap_or("".to_string());
28    for s2 in i {
29        s.push_str(".");
30        s.push_str(&s2.to_string()[..]);
31    }
32    s
33}
34
35pub fn format_module(stmts: &[Statement]) -> String {
36    let mut s = "".to_string();
37    for stmt in stmts {
38        s.push_str(&format_statement(0, &stmt)[..])
39    }
40    s
41}
42
43fn push_indent(indent: usize, s: &mut String) {
44    for _ in 0..indent {
45        s.push_str(" ")
46    }
47}
48
49fn format_statement(indent: usize, stmt: &Statement) -> String {
50    let mut s = "".to_string();
51    push_indent(indent, &mut s);
52    match *stmt {
53        Statement::Pass => s.push_str("pass\n"),
54        Statement::Del(ref exprs) => {
55            s.push_str("del ");
56            s.push_str(&comma_join(exprs.iter().map(format_expr)));
57            s.push_str("\n");
58        }
59        Statement::Break => s.push_str("break\n"),
60        Statement::Continue => s.push_str("continue\n"),
61        Statement::Return(ref exprs) => {
62            s.push_str("return ");
63            s.push_str(&comma_join(exprs.iter().map(format_expr)));
64            s.push_str("\n");
65        }
66        Statement::RaiseExcFrom(ref exc, ref from_exc) => {
67            s.push_str("raise ");
68            s.push_str(&format_expr(exc));
69            s.push_str(" from ");
70            s.push_str(&format_expr(from_exc));
71            s.push_str("\n");
72        }
73        Statement::RaiseExc(ref exc) => {
74            s.push_str("raise ");
75            s.push_str(&format_expr(exc));
76            s.push_str("\n");
77        }
78        Statement::Raise => {
79            s.push_str("raise\n");
80        }
81        Statement::Global(ref names) => {
82            s.push_str("global ");
83            s.push_str(&comma_join(names));
84            s.push_str("\n");
85        }
86        Statement::Nonlocal(ref names) => {
87            s.push_str("nonlocal ");
88            s.push_str(&comma_join(names));
89            s.push_str("\n");
90        }
91        Statement::Assert(ref expr, ref msg) => {
92            s.push_str("assert ");
93            s.push_str(&format_expr(expr));
94            if let &Some(ref msg) = msg {
95                s.push_str(", ");
96                s.push_str(&format_expr(msg));
97            }
98            s.push_str("\n");
99        }
100        Statement::Import(ref imp) => {
101            s.push_str(&format_import(imp));
102            s.push_str("\n");
103        }
104        Statement::Expressions(ref exprs) => {
105            s.push_str(&comma_join(exprs.iter().map(format_expr)));
106            s.push_str("\n");
107        }
108        Statement::Assignment(ref lhs, ref rhs) => {
109            s.push_str(&comma_join(lhs.iter().map(format_expr)));
110            for part in rhs {
111                s.push_str(" = ");
112                s.push_str(&comma_join(part.iter().map(format_expr)));
113            }
114            s.push_str("\n");
115        }
116        Statement::TypeAnnotation(ref lhs, ref typed) => {
117            s.push_str(&format!(
118                "{}: {}\n",
119                comma_join(lhs.iter().map(format_expr)),
120                format_expr(typed)
121            ));
122        }
123        Statement::TypedAssignment(ref lhs, ref typed, ref rhs) => {
124            s.push_str(&format!(
125                "{}:{} = {}\n",
126                comma_join(lhs.iter().map(format_expr)),
127                format_expr(typed),
128                comma_join(rhs.iter().map(format_expr))
129            ));
130        }
131        Statement::AugmentedAssignment(ref lhs, op, ref rhs) => {
132            s.push_str(&format!(
133                "{} {} {}\n",
134                comma_join(lhs.iter().map(format_expr)),
135                op,
136                comma_join(rhs.iter().map(format_expr))
137            ));
138        }
139        Statement::Compound(ref stmt) => s.push_str(&format_compound_statement(indent, stmt)),
140    }
141    s
142}
143
144fn format_compound_statement(indent: usize, stmt: &CompoundStatement) -> String {
145    match *stmt {
146        CompoundStatement::If(ref cond_blocks, ref else_block) => {
147            let mut s = String::new();
148            let mut first = true;
149            for &(ref cond, ref block) in cond_blocks {
150                if first {
151                    s.push_str("if ");
152                    s.push_str(&format_expr(cond));
153                    s.push_str(":\n");
154                    s.push_str(&format_block(indent + 4, block));
155                    first = false;
156                } else {
157                    push_indent(indent, &mut s);
158                    s.push_str("elif ");
159                    s.push_str(&format_expr(cond));
160                    s.push_str(":\n");
161                    s.push_str(&format_block(indent + 4, block));
162                }
163            }
164            if let &Some(ref block) = else_block {
165                push_indent(indent, &mut s);
166                s.push_str("else:\n");
167                s.push_str(&format_block(indent + 4, block));
168            }
169            s
170        }
171        CompoundStatement::For {
172            async,
173            ref item,
174            ref iterator,
175            ref for_block,
176            ref else_block,
177        } => {
178            let mut s = String::new();
179            if async {
180                s.push_str("async ");
181            }
182            s.push_str("for ");
183            s.push_str(&comma_join(item.iter().map(format_expr)));
184            s.push_str(" in ");
185            s.push_str(&comma_join(iterator.iter().map(format_expr)));
186            s.push_str(":\n");
187            s.push_str(&format_block(indent + 4, for_block));
188
189            if let &Some(ref block) = else_block {
190                push_indent(indent, &mut s);
191                s.push_str("else:\n");
192                s.push_str(&format_block(indent + 4, block));
193            }
194            s
195        }
196        CompoundStatement::While(ref cond, ref block, ref else_block) => {
197            let mut s = String::new();
198            s.push_str("while ");
199            s.push_str(&format_expr(cond));
200            s.push_str(":\n");
201            s.push_str(&format_block(indent + 4, block));
202
203            if let &Some(ref block) = else_block {
204                push_indent(indent, &mut s);
205                s.push_str("else:\n");
206                s.push_str(&format_block(indent + 4, block));
207            }
208            s
209        }
210        CompoundStatement::Try(Try {
211            ref try_block,
212            ref except_clauses,
213            ref last_except,
214            ref else_block,
215            ref finally_block,
216        }) => {
217            let mut s = String::new();
218
219            s.push_str("try:\n");
220            s.push_str(&format_block(indent + 4, try_block));
221
222            for &(ref guard, ref name, ref block) in except_clauses {
223                push_indent(indent, &mut s);
224                s.push_str("except ");
225                s.push_str(&format_expr(guard));
226                if let &Some(ref name) = name {
227                    s.push_str(" as ");
228                    s.push_str(name);
229                }
230                s.push_str(":\n");
231                s.push_str(&format_block(indent + 4, block));
232            }
233            if last_except.len() > 0 {
234                push_indent(indent, &mut s);
235                s.push_str("except:\n");
236                s.push_str(&format_block(indent + 4, last_except));
237            }
238            if else_block.len() > 0 {
239                push_indent(indent, &mut s);
240                s.push_str("else:\n");
241                s.push_str(&format_block(indent + 4, else_block));
242            }
243            if finally_block.len() > 0 {
244                push_indent(indent, &mut s);
245                s.push_str("finally:\n");
246                s.push_str(&format_block(indent + 4, finally_block));
247            }
248            s
249        }
250        CompoundStatement::With(ref contexts, ref block) => {
251            let mut s = String::new();
252
253            s.push_str("with ");
254            assert!(contexts.len() > 0);
255            let mut first = true;
256            for &(ref ctx, ref as_what) in contexts {
257                if first {
258                    first = false;
259                } else {
260                    s.push_str(", ");
261                }
262                s.push_str(&format_expr(ctx));
263                if let &Some(ref e) = as_what {
264                    s.push_str(" as ");
265                    s.push_str(&format_expr(e));
266                }
267            }
268            s.push_str(":\n");
269            s.push_str(&format_block(indent + 4, block));
270            s
271        }
272        CompoundStatement::Funcdef(ref funcdef) => format_funcdef(indent, funcdef),
273        CompoundStatement::Classdef(ref classdef) => format_classdef(indent, classdef),
274    }
275}
276
277fn format_decorators(indent: usize, decorators: &Vec<Decorator>) -> String {
278    let mut s = String::new();
279    for &Decorator { ref name, ref args } in decorators {
280        push_indent(indent, &mut s);
281        s.push_str("@");
282        s.push_str(&dot_join(name));
283        if let &Some(ref arglist) = args {
284            s.push_str("(");
285            s.push_str(&format_args(arglist));
286            s.push_str(")");
287        }
288        s.push_str("\n");
289    }
290    s
291}
292
293fn format_funcdef(indent: usize, funcdef: &Funcdef) -> String {
294    let &Funcdef {
295        async,
296        ref decorators,
297        ref name,
298        ref parameters,
299        ref return_type,
300        ref code,
301    } = funcdef;
302    let mut s = "\n".to_string();
303    s.push_str(&format_decorators(indent, decorators));
304    push_indent(indent, &mut s);
305    if async {
306        s.push_str("async ");
307    }
308    s.push_str("def ");
309    s.push_str(name);
310    s.push_str("(");
311    s.push_str(&format_typed_params(parameters));
312    s.push_str(")");
313    if let &Some(ref ret) = return_type {
314        s.push_str(" -> ");
315        s.push_str(&format_expr(ret));
316    }
317    s.push_str(":\n");
318    s.push_str(&format_block(indent + 4, code));
319    s.push_str("\n");
320    s
321}
322
323fn format_classdef(indent: usize, classdef: &Classdef) -> String {
324    let &Classdef {
325        ref decorators,
326        ref name,
327        ref arguments,
328        ref code,
329    } = classdef;
330    let mut s = "\n".to_string();
331    s.push_str(&format_decorators(indent, decorators));
332    push_indent(indent, &mut s);
333    s.push_str("class ");
334    s.push_str(name);
335    s.push_str("(");
336    s.push_str(&format_args(arguments));
337    s.push_str(")");
338    s.push_str(":\n");
339    s.push_str(&format_block(indent + 4, code));
340    s.push_str("\n");
341    s
342}
343
344fn format_block(indent: usize, stmts: &Vec<Statement>) -> String {
345    let mut s = String::new();
346    for stmt in stmts {
347        s.push_str(&format_statement(indent, stmt));
348    }
349    s
350}
351
352fn format_dictitem(si: &DictItem) -> String {
353    match *si {
354        DictItem::Unique(ref e1, ref e2) => format!("{}:{}", format_expr(e1), format_expr(e2)),
355        DictItem::Star(ref e) => format!("**{}", format_expr(e)),
356    }
357}
358
359fn format_setitem(si: &SetItem) -> String {
360    match *si {
361        SetItem::Unique(ref e) => format_expr(e),
362        SetItem::Star(ref e) => format!("*{}", format_expr(e)),
363    }
364}
365
366fn format_args(args: &Vec<Argument>) -> String {
367    let mut s = String::new();
368    s.push_str(&comma_join(args.iter().map(|arg| match *arg {
369        Argument::Positional(ref e) => format_expr(e),
370        Argument::Starargs(ref e) => format!("*{}", format_expr(e)),
371        Argument::Keyword(ref n, ref e) => format!("{}={}", n, format_expr(e)),
372        Argument::Kwargs(ref e) => format!("**{}", format_expr(e)),
373    })));
374    s
375}
376
377fn format_typed_params(param: &TypedArgsList) -> String {
378    let TypedArgsList {
379        ref posonly_args,
380        ref args,
381        ref star_args,
382        ref keyword_args,
383        ref star_kwargs,
384    } = *param;
385    let mut chunks = Vec::new();
386
387    if posonly_args.len() > 0 {
388        chunks.extend(posonly_args.iter().map(format_typed_param));
389        chunks.push("/".to_string());
390    }
391
392    chunks.extend(args.iter().map(format_typed_param));
393
394    match *star_args {
395        StarParams::No => (),
396        StarParams::Anonymous => chunks.push("*".to_string()),
397        StarParams::Named((ref name, None)) => chunks.push(format!("*{}", name)),
398        StarParams::Named((ref name, Some(ref typed))) => {
399            chunks.push(format!("*{}:{}", name, format_expr(typed)))
400        }
401    }
402
403    chunks.extend(keyword_args.iter().map(format_typed_param));
404
405    if let &Some((ref name, ref typed)) = star_kwargs {
406        if let &Some(ref typed) = typed {
407            chunks.push(format!("**{}:{}", name, format_expr(typed)))
408        } else {
409            chunks.push(format!("**{}", name));
410        }
411    }
412
413    comma_join(chunks)
414}
415
416fn format_typed_param(param: &(Name, Option<Expression>, Option<Expression>)) -> String {
417    let &(ref name, ref typed, ref value) = param;
418    let mut s = name.to_string();
419    if let &Some(ref typed) = typed {
420        s.push_str(":");
421        s.push_str(&format_expr(typed));
422    }
423    if let &Some(ref value) = value {
424        s.push_str("=");
425        s.push_str(&format_expr(value));
426    }
427    s
428}
429
430fn format_untyped_params(param: &UntypedArgsList) -> String {
431    let UntypedArgsList {
432        ref posonly_args,
433        ref args,
434        ref star_args,
435        ref keyword_args,
436        ref star_kwargs,
437    } = *param;
438
439    let mut chunks = Vec::new();
440
441    if posonly_args.len() > 0 {
442        chunks.extend(posonly_args.iter().map(format_untyped_param));
443        chunks.push("/".to_string());
444    }
445
446    chunks.extend(args.iter().map(format_untyped_param));
447
448    match *star_args {
449        StarParams::No => (),
450        StarParams::Anonymous => chunks.push("*".to_string()),
451        StarParams::Named(ref name) => chunks.push(format!("*{}", name)),
452    }
453
454    chunks.extend(keyword_args.iter().map(format_untyped_param));
455
456    if let &Some(ref name) = star_kwargs {
457        chunks.push(format!("**{}", name));
458    }
459
460    comma_join(&chunks)
461}
462
463fn format_untyped_param(param: &(Name, Option<Expression>)) -> String {
464    let &(ref name, ref value) = param;
465    let mut s = name.to_string();
466    if let &Some(ref value) = value {
467        s.push_str("=");
468        s.push_str(&format_expr(value));
469    }
470    s
471}
472
473fn format_subscript(sub: &Subscript) -> String {
474    match *sub {
475        Subscript::Simple(ref e) => format_expr(e),
476        Subscript::Double(ref e1, ref e2) => format!(
477            "{}:{}",
478            e1.clone().map(|e| format_expr(&e)).unwrap_or_default(),
479            e2.clone().map(|e| format_expr(&e)).unwrap_or_default(),
480        ),
481        Subscript::Triple(ref e1, ref e2, ref e3) => format!(
482            "{}:{}:{}",
483            e1.clone().map(|e| format_expr(&e)).unwrap_or_default(),
484            e2.clone().map(|e| format_expr(&e)).unwrap_or_default(),
485            e3.clone().map(|e| format_expr(&e)).unwrap_or_default(),
486        ),
487    }
488}
489
490fn format_comp(comp: &ComprehensionChunk) -> String {
491    match *comp {
492        ComprehensionChunk::If { ref cond } => format!("if ({})", format_expr(cond)),
493        ComprehensionChunk::For {
494            async,
495            ref item,
496            ref iterator,
497        } => format!(
498            "{}for {} in {}",
499            if async { "async " } else { "" },
500            comma_join(item.iter().map(format_expr)),
501            format_expr(iterator)
502        ),
503    }
504}
505
506fn format_float(n: f64) -> String {
507    let mut s = n.to_string();
508    if s.find('.').is_none() {
509        s.push_str(".");
510    }
511    s
512}
513
514#[cfg(feature = "wtf8")]
515fn format_string(v: &Vec<PyString>) -> String {
516    space_join(v.iter().map(
517        |&PyString {
518             ref prefix,
519             ref content,
520         }| {
521            format!(
522                "{}\"{}\"",
523                prefix.to_ascii_lowercase().replace("r", ""),
524                content
525                    .code_points()
526                    .map(|c| match c.to_u32() {
527                        0xd => "\\r".to_string(),
528                        0xa => "\\n".to_string(),
529                        0x9 => "\\t".to_string(),
530                        0x5c => "\\\\".to_string(),
531                        0x22 => "\\\"".to_string(),
532                        0x20..=0x7e => c.to_char().unwrap().to_string(), // unwrap can't panic
533                        0x00..=0x1f | 0x7f | 0x80..=0xff => format!("\\x{:02x}", c.to_u32()),
534                        0x100..=0xffff => format!("\\u{:04x}", c.to_u32()),
535                        0x10000..=0x10ffff => format!("\\U{:08x}", c.to_u32()),
536                        _ => unreachable!(),
537                    })
538                    .collect::<Vec<_>>()[..]
539                    .concat()
540            )
541        },
542    ))
543}
544
545#[cfg(not(feature = "wtf8"))]
546fn format_string(v: &Vec<PyString>) -> String {
547    space_join(v.iter().map(
548        |&PyString {
549             ref prefix,
550             ref content,
551         }| {
552            format!(
553                "{}\"{}\"",
554                prefix.to_ascii_lowercase().replace("r", ""),
555                content
556                    .chars()
557                    .map(|c| match c {
558                        '\r' => "\\r".to_string(),
559                        '\n' => "\\n".to_string(),
560                        '\t' => "\\t".to_string(),
561                        '\\' => "\\\\".to_string(),
562                        '"' => "\\\"".to_string(),
563                        '\x20'...'\x7e' => c.to_string(),
564                        '\x00'...'\x1f' | '\x7f' | '\u{80}'...'\u{ff}' =>
565                            format!("\\x{:02x}", c as u8),
566                        '\u{100}'...'\u{ffff}' => format!("\\u{:04x}", c as u16),
567                        '\u{10000}'...'\u{10ffff}' => format!("\\U{:08x}", c as u32),
568                        _ => unreachable!(),
569                    })
570                    .collect::<Vec<_>>()[..]
571                    .concat()
572            )
573        },
574    ))
575}
576
577fn format_expr(e: &Expression) -> String {
578    match *e {
579        Expression::Ellipsis => "...".to_string(),
580        Expression::None => "None".to_string(),
581        Expression::True => "True".to_string(),
582        Expression::False => "False".to_string(),
583        Expression::Name(ref n) => n.to_string(),
584        Expression::Int(ref n) => n.to_string(),
585        Expression::ImaginaryInt(ref n) => format!("{}j", n),
586        Expression::Float(ref n) => format_float(*n),
587        Expression::ImaginaryFloat(ref n) => format!("{}j", format_float(*n)),
588        Expression::String(ref v) => format_string(v),
589        Expression::Bytes(ref content) => format!(
590            "b\"{}\"",
591            content
592                .iter()
593                .map(|b| match *b {
594                    b'\r' => "\\r".to_string(),
595                    b'\n' => "\\n".to_string(),
596                    b'\t' => "\\t".to_string(),
597                    b'\\' => "\\\\".to_string(),
598                    b'"' => "\\\"".to_string(),
599                    0x20..=0x7e => (*b as char).to_string(),
600                    0x00..=0x1f | 0x7f | 0x80..=0xff => format!("\\x{:02x}", b),
601                })
602                .collect::<Vec<_>>()[..]
603                .concat()
604        ),
605
606        Expression::DictLiteral(ref v) => {
607            format!("{{{}}}", comma_join(v.iter().map(format_dictitem)))
608        }
609        Expression::SetLiteral(ref v) => {
610            format!("{{{}}}", comma_join(v.iter().map(format_setitem)))
611        }
612        Expression::ListLiteral(ref v) => format!("[{}]", comma_join(v.iter().map(format_setitem))),
613        Expression::TupleLiteral(ref v) => match v.len() {
614            0 => "()".to_string(),
615            1 => format!("({},)", format_setitem(&v[0])),
616            _ => format!("({})", comma_join(v.iter().map(format_setitem))),
617        },
618
619        Expression::DictComp(ref e, ref comp) => format!(
620            "{{{} {}}}",
621            format_dictitem(e),
622            space_join(comp.iter().map(format_comp))
623        ),
624        Expression::SetComp(ref e, ref comp) => format!(
625            "{{{} {}}}",
626            format_setitem(e),
627            space_join(comp.iter().map(format_comp))
628        ),
629        Expression::ListComp(ref e, ref comp) => format!(
630            "[{} {}]",
631            format_setitem(e),
632            space_join(comp.iter().map(format_comp))
633        ),
634        Expression::Generator(ref e, ref comp) => format!(
635            "({} {})",
636            format_setitem(e),
637            space_join(comp.iter().map(format_comp))
638        ),
639        Expression::Await(ref e) => format!("await {}", format_expr(e)),
640
641        Expression::Call(ref e, ref args) => match **e {
642            Expression::Name(_)
643            | Expression::DictComp(_, _)
644            | Expression::SetComp(_, _)
645            | Expression::ListComp(_, _)
646            | Expression::Generator(_, _)
647            | Expression::DictLiteral(_)
648            | Expression::SetLiteral(_)
649            | Expression::ListLiteral(_)
650            | Expression::TupleLiteral(_)
651            | Expression::Attribute(_, _)
652            | Expression::Call(_, _) => format!("{}({})", format_expr(e), format_args(args)),
653            _ => format!("({})({})", format_expr(e), format_args(args)),
654        },
655        Expression::Subscript(ref e, ref sub) => format!(
656            "({})[{}]",
657            format_expr(e),
658            comma_join(sub.iter().map(format_subscript))
659        ),
660        Expression::Attribute(ref e, ref n) => match **e {
661            Expression::Name(_)
662            | Expression::DictComp(_, _)
663            | Expression::SetComp(_, _)
664            | Expression::ListComp(_, _)
665            | Expression::Generator(_, _)
666            | Expression::DictLiteral(_)
667            | Expression::SetLiteral(_)
668            | Expression::ListLiteral(_)
669            | Expression::TupleLiteral(_)
670            | Expression::Attribute(_, _)
671            | Expression::Call(_, _) => format!("{}.{}", format_expr(e), n),
672            _ => format!("({}).{}", format_expr(e), n),
673        },
674        Expression::Uop(op, ref e) => format!("{}({})", op, format_expr(e)),
675        Expression::Bop(op, ref e1, ref e2) => {
676            let f = |e: &_| match *e {
677                Expression::Ellipsis
678                | Expression::None
679                | Expression::True
680                | Expression::False
681                | Expression::Int(_)
682                | Expression::ImaginaryInt(_)
683                | Expression::ImaginaryFloat(_)
684                | Expression::Float(_)
685                | Expression::String(_)
686                | Expression::Bytes(_)
687                | Expression::Name(_)
688                | Expression::DictComp(_, _)
689                | Expression::SetComp(_, _)
690                | Expression::ListComp(_, _)
691                | Expression::Generator(_, _)
692                | Expression::DictLiteral(_)
693                | Expression::SetLiteral(_)
694                | Expression::ListLiteral(_)
695                | Expression::TupleLiteral(_)
696                | Expression::Attribute(_, _)
697                | Expression::Call(_, _) => format!("{}", format_expr(e)),
698                _ => format!("({})", format_expr(e)),
699            };
700            format!("{}{}{}", f(e1), op, f(e2))
701        }
702        Expression::MultiBop(ref first, ref rest) => {
703            let mut s = String::new();
704            s.push_str("(");
705            s.push_str(&format_expr(first));
706            s.push_str(")");
707            for &(op, ref e) in rest {
708                s.push_str(" ");
709                s.push_str(&op.to_string());
710                s.push_str(" (");
711                s.push_str(&format_expr(e));
712                s.push_str(")");
713            }
714            s
715        }
716        Expression::Ternary(ref e1, ref e2, ref e3) => format!(
717            "({}) if ({}) else ({})",
718            format_expr(e1),
719            format_expr(e2),
720            format_expr(e3)
721        ),
722        Expression::Star(ref e) => format!("*{}", format_expr(e)),
723
724        // https://mail.python.org/pipermail/python-list/2013-August/653288.html
725        Expression::Yield(ref items) => {
726            format!("(yield {})", comma_join(items.iter().map(format_expr)))
727        }
728        Expression::YieldFrom(ref iterable) => format!("(yield from {})", format_expr(iterable)),
729
730        Expression::Lambdef(ref params, ref body) => format!(
731            "lambda {}: {}",
732            format_untyped_params(params),
733            format_expr(body)
734        ),
735        Expression::Named(ref name, ref expr) => {
736            format!("{} := ({})", format_expr(name), format_expr(expr),)
737        }
738    }
739}
740
741fn format_dotted_name(path: &[String]) -> String {
742    let mut s = "".to_string();
743    let mut first = true;
744    for part in path {
745        if first {
746            first = false;
747        } else {
748            s.push_str(".");
749        }
750        s.push_str(part);
751    }
752    s
753}
754
755fn format_import(imp: &Import) -> String {
756    let mut s = "".to_string();
757    match *imp {
758        Import::ImportFrom {
759            leading_dots,
760            ref path,
761            ref names,
762        } => {
763            s.push_str("from ");
764            for _ in 0..leading_dots {
765                s.push_str(".");
766            }
767            s.push_str(&format_dotted_name(path));
768            s.push_str(" import ");
769            s.push_str(&comma_join(names.iter().map(|&(ref name, ref as_name)| {
770                let mut s2 = String::new();
771                s2.push_str(name);
772                if let &Some(ref as_name) = as_name {
773                    s2.push_str(" as ");
774                    s2.push_str(as_name);
775                }
776                s2
777            })));
778        }
779        Import::ImportStarFrom {
780            leading_dots,
781            ref path,
782        } => {
783            s.push_str("from ");
784            for _ in 0..leading_dots {
785                s.push_str(".");
786            }
787            s.push_str(&format_dotted_name(path));
788            s.push_str(" import *");
789        }
790        Import::Import { ref names } => {
791            s.push_str("import ");
792            s.push_str(&comma_join(names.iter().map(|&(ref name, ref as_name)| {
793                let mut s2 = String::new();
794                s2.push_str(&format_dotted_name(name));
795                if let &Some(ref as_name) = as_name {
796                    s2.push_str(" as ");
797                    s2.push_str(as_name);
798                }
799                s2
800            })));
801        }
802    }
803    s
804}
805
806#[cfg(test)]
807mod tests {
808    use super::*;
809
810    #[test]
811    fn test_ternary_in_comp_if() {
812        let e = Expression::ListComp(
813            Box::new(SetItem::Unique(Expression::Name("a".to_string()))),
814            vec![
815                ComprehensionChunk::For {
816                    async: false,
817                    item: vec![Expression::Name("a".to_string())],
818                    iterator: Expression::Name("L".to_string()),
819                },
820                ComprehensionChunk::If {
821                    cond: Expression::Ternary(
822                        Box::new(Expression::Call(
823                            Box::new(Expression::Name("f".to_string())),
824                            vec![Argument::Positional(Expression::Name("a".to_string()))],
825                        )),
826                        Box::new(Expression::Name("a".to_string())),
827                        Box::new(Expression::None),
828                    ),
829                },
830            ],
831        );
832        assert_eq!(
833            &format_expr(&e),
834            "[a for a in L if ((f(a)) if (a) else (None))]"
835        );
836    }
837
838    #[test]
839    fn test_namedexpr() {
840        let e = Expression::Named(
841            Box::new(Expression::Name("foo".to_string())),
842            Box::new(Expression::Name("bar".to_string())),
843        );
844        assert_eq!(&format_expr(&e), "foo := (bar)");
845    }
846}