blueprint_starlark_syntax/syntax/
ast.rs

1/*
2 * Copyright 2018 The Starlark in Rust Authors.
3 * Copyright (c) Facebook, Inc. and its affiliates.
4 *
5 * Licensed under the Apache License, Version 2.0 (the "License");
6 * you may not use this file except in compliance with the License.
7 * You may obtain a copy of the License at
8 *
9 *     https://www.apache.org/licenses/LICENSE-2.0
10 *
11 * Unless required by applicable law or agreed to in writing, software
12 * distributed under the License is distributed on an "AS IS" BASIS,
13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 * See the License for the specific language governing permissions and
15 * limitations under the License.
16 */
17
18//! AST for parsed starlark files.
19
20use std::fmt;
21use std::fmt::Debug;
22use std::fmt::Display;
23use std::fmt::Formatter;
24
25use blueprint_allocative::Allocative;
26use blueprint_dupe::Dupe;
27
28use crate::codemap::Pos;
29use crate::codemap::Span;
30use crate::codemap::Spanned;
31use crate::lexer::TokenInt;
32
33/// Payload types attached to AST nodes.
34pub trait AstPayload: Debug {
35    // TODO(nga): we don't really need `Clone` for any payload.
36    type LoadPayload: Debug + Clone;
37    type IdentPayload: Debug + Clone;
38    type IdentAssignPayload: Debug + Clone;
39    type DefPayload: Debug + Clone;
40    type TypeExprPayload: Debug + Clone;
41}
42
43/// Default implementation of payload, which attaches `()` to nodes.
44/// This payload is returned with AST by parser.
45#[derive(Debug, Copy, Clone, Dupe, Eq, PartialEq)]
46pub struct AstNoPayload;
47impl AstPayload for AstNoPayload {
48    type LoadPayload = ();
49    type IdentPayload = ();
50    type IdentAssignPayload = ();
51    type DefPayload = ();
52    type TypeExprPayload = ();
53}
54
55/// `,` token.
56#[derive(Copy, Clone, Dupe, Debug)]
57pub struct Comma;
58
59pub type Expr = ExprP<AstNoPayload>;
60pub type TypeExpr = TypeExprP<AstNoPayload>;
61pub type AssignTarget = AssignTargetP<AstNoPayload>;
62pub type AssignIdent = AssignIdentP<AstNoPayload>;
63pub type Ident = IdentP<AstNoPayload>;
64pub type Clause = ClauseP<AstNoPayload>;
65pub type ForClause = ForClauseP<AstNoPayload>;
66pub type Argument = ArgumentP<AstNoPayload>;
67pub type Parameter = ParameterP<AstNoPayload>;
68pub type Load = LoadP<AstNoPayload>;
69pub type Stmt = StmtP<AstNoPayload>;
70
71// Boxed types used for storing information from the parsing will be used
72// especially for the location of the AST item
73pub type AstExprP<P> = Spanned<ExprP<P>>;
74pub type AstTypeExprP<P> = Spanned<TypeExprP<P>>;
75pub type AstAssignTargetP<P> = Spanned<AssignTargetP<P>>;
76pub type AstAssignIdentP<P> = Spanned<AssignIdentP<P>>;
77pub type AstIdentP<P> = Spanned<IdentP<P>>;
78pub type AstArgumentP<P> = Spanned<ArgumentP<P>>;
79pub type AstParameterP<P> = Spanned<ParameterP<P>>;
80pub type AstStmtP<P> = Spanned<StmtP<P>>;
81pub type AstFStringP<P> = Spanned<FStringP<P>>;
82
83pub type AstExpr = AstExprP<AstNoPayload>;
84pub type AstTypeExpr = AstTypeExprP<AstNoPayload>;
85pub type AstAssignTarget = AstAssignTargetP<AstNoPayload>;
86pub type AstAssignIdent = AstAssignIdentP<AstNoPayload>;
87pub type AstIdent = AstIdentP<AstNoPayload>;
88pub type AstArgument = AstArgumentP<AstNoPayload>;
89pub type AstString = Spanned<String>;
90pub type AstParameter = AstParameterP<AstNoPayload>;
91pub type AstInt = Spanned<TokenInt>;
92pub type AstFloat = Spanned<f64>;
93pub type AstFString = AstFStringP<AstNoPayload>;
94pub type AstStmt = AstStmtP<AstNoPayload>;
95
96// A trait rather than a function to allow .ast() chaining in the parser.
97pub trait ToAst: Sized {
98    fn ast(self, begin: usize, end: usize) -> Spanned<Self> {
99        Spanned {
100            span: Span::new(Pos::new(begin as u32), Pos::new(end as u32)),
101            node: self,
102        }
103    }
104}
105
106impl<T> ToAst for T {}
107
108#[derive(Debug, Clone)]
109pub enum ArgumentP<P: AstPayload> {
110    Positional(AstExprP<P>),
111    Named(AstString, AstExprP<P>),
112    Args(AstExprP<P>),
113    KwArgs(AstExprP<P>),
114}
115
116#[derive(Debug, Clone)]
117pub enum ParameterP<P: AstPayload> {
118    /// `/` marker.
119    Slash,
120    Normal(
121        /// Name.
122        AstAssignIdentP<P>,
123        /// Type.
124        Option<Box<AstTypeExprP<P>>>,
125        /// Default value.
126        Option<Box<AstExprP<P>>>,
127    ),
128    /// `*` marker.
129    NoArgs,
130    Args(AstAssignIdentP<P>, Option<Box<AstTypeExprP<P>>>),
131    KwArgs(AstAssignIdentP<P>, Option<Box<AstTypeExprP<P>>>),
132}
133
134impl<P: AstPayload> ParameterP<P> {
135    pub fn ident(&self) -> Option<&AstAssignIdentP<P>> {
136        match self {
137            ParameterP::Normal(x, _, _) | ParameterP::Args(x, _) | ParameterP::KwArgs(x, _) => {
138                Some(x)
139            }
140            ParameterP::NoArgs | ParameterP::Slash => None,
141        }
142    }
143}
144
145#[derive(Debug, Clone)]
146pub enum AstLiteral {
147    Int(AstInt),
148    Float(AstFloat),
149    String(AstString),
150    Ellipsis,
151}
152
153#[derive(Debug, Clone)]
154pub struct LambdaP<P: AstPayload> {
155    pub params: Vec<AstParameterP<P>>,
156    pub body: Box<AstExprP<P>>,
157    pub payload: P::DefPayload,
158}
159
160impl<P: AstPayload> LambdaP<P> {
161    pub fn signature_span(&self) -> Span {
162        self.params
163            .iter()
164            .map(|p| p.span)
165            .reduce(|a, b| a.merge(b))
166            .unwrap_or(
167                // TODO(nga): this is not correct span.
168                self.body.span,
169            )
170    }
171}
172
173#[derive(Debug, Clone)]
174pub struct CallArgsP<P: AstPayload> {
175    pub args: Vec<AstArgumentP<P>>,
176}
177
178#[derive(Debug, Clone)]
179pub enum ExprP<P: AstPayload> {
180    Tuple(Vec<AstExprP<P>>),
181    Dot(Box<AstExprP<P>>, AstString),
182    Call(Box<AstExprP<P>>, CallArgsP<P>),
183    Index(Box<(AstExprP<P>, AstExprP<P>)>),
184    Index2(Box<(AstExprP<P>, AstExprP<P>, AstExprP<P>)>),
185    Slice(
186        Box<AstExprP<P>>,
187        Option<Box<AstExprP<P>>>,
188        Option<Box<AstExprP<P>>>,
189        Option<Box<AstExprP<P>>>,
190    ),
191    Identifier(AstIdentP<P>),
192    Lambda(LambdaP<P>),
193    Literal(AstLiteral),
194    Not(Box<AstExprP<P>>),
195    Minus(Box<AstExprP<P>>),
196    Plus(Box<AstExprP<P>>),
197    BitNot(Box<AstExprP<P>>),
198    Op(Box<AstExprP<P>>, BinOp, Box<AstExprP<P>>),
199    If(Box<(AstExprP<P>, AstExprP<P>, AstExprP<P>)>), // Order: condition, v1, v2 <=> v1 if condition else v2
200    List(Vec<AstExprP<P>>),
201    Dict(Vec<(AstExprP<P>, AstExprP<P>)>),
202    Set(Vec<AstExprP<P>>),
203    ListComprehension(Box<AstExprP<P>>, Box<ForClauseP<P>>, Vec<ClauseP<P>>),
204    SetComprehension(Box<AstExprP<P>>, Box<ForClauseP<P>>, Vec<ClauseP<P>>),
205    DictComprehension(
206        Box<(AstExprP<P>, AstExprP<P>)>,
207        Box<ForClauseP<P>>,
208        Vec<ClauseP<P>>,
209    ),
210    FString(AstFStringP<P>),
211}
212
213/// Restricted expression at type position.
214#[derive(Debug, Clone)]
215pub struct TypeExprP<P: AstPayload> {
216    /// Currently it is an expr.
217    /// Planning to restrict it.
218    /// [Context](https://fb.workplace.com/groups/buck2eng/posts/3196541547309990).
219    pub expr: AstExprP<P>,
220    pub payload: P::TypeExprPayload,
221}
222
223/// In some places e.g. AssignModify, the Tuple case is not allowed.
224#[derive(Debug, Clone)]
225pub enum AssignTargetP<P: AstPayload> {
226    // We use Tuple for both Tuple and List,
227    // as these have the same semantics in Starlark.
228    Tuple(Vec<AstAssignTargetP<P>>),
229    Index(Box<(AstExprP<P>, AstExprP<P>)>),
230    Dot(Box<AstExprP<P>>, AstString),
231    Identifier(AstAssignIdentP<P>),
232}
233
234/// `x: t = y`.
235#[derive(Debug, Clone)]
236pub struct AssignP<P: AstPayload> {
237    pub lhs: AstAssignTargetP<P>,
238    pub ty: Option<AstTypeExprP<P>>,
239    pub rhs: AstExprP<P>,
240}
241
242/// Identifier in assign position.
243#[derive(Debug, Eq, PartialEq, Clone)]
244pub struct AssignIdentP<P: AstPayload> {
245    pub ident: String,
246    pub payload: P::IdentAssignPayload,
247}
248
249/// Identifier in read position, e. g. `foo` in `[foo.bar]`.
250/// `foo` in `foo = 1` or `bar.foo` are **not** represented by this type.
251#[derive(Debug, Eq, PartialEq, Clone)]
252pub struct IdentP<P: AstPayload> {
253    pub ident: String,
254    pub payload: P::IdentPayload,
255}
256
257/// Argument of `load` statement.
258#[derive(Debug, Clone)]
259pub struct LoadArgP<P: AstPayload> {
260    /// `x in `x="y"`.
261    pub local: AstAssignIdentP<P>,
262    /// `"y" in `x="y"`.
263    pub their: AstString,
264    /// Trailing comma.
265    pub comma: Option<Spanned<Comma>>,
266}
267
268impl<P: AstPayload> LoadArgP<P> {
269    pub fn span(&self) -> Span {
270        self.local.span.merge(self.their.span)
271    }
272
273    pub fn span_with_trailing_comma(&self) -> Span {
274        if let Some(comma) = &self.comma {
275            self.span().merge(comma.span)
276        } else {
277            self.span()
278        }
279    }
280}
281
282/// `load` statement.
283#[derive(Debug, Clone)]
284pub struct LoadP<P: AstPayload> {
285    pub module: AstString,
286    pub args: Vec<LoadArgP<P>>,
287    pub payload: P::LoadPayload,
288}
289
290#[derive(Debug, Clone)]
291pub struct ForClauseP<P: AstPayload> {
292    pub var: AstAssignTargetP<P>,
293    pub over: AstExprP<P>,
294}
295
296#[derive(Debug, Clone)]
297pub enum ClauseP<P: AstPayload> {
298    For(ForClauseP<P>),
299    If(AstExprP<P>),
300}
301
302#[derive(Debug, Clone, Copy, Dupe, Eq, PartialEq)]
303pub enum BinOp {
304    Or,
305    And,
306    Equal,
307    NotEqual,
308    Less,
309    Greater,
310    LessOrEqual,
311    GreaterOrEqual,
312    In,
313    NotIn,
314    Subtract,
315    Add,
316    Multiply,
317    Percent,
318    Divide,
319    FloorDivide,
320    BitAnd,
321    BitOr,
322    BitXor,
323    LeftShift,
324    RightShift,
325}
326
327#[derive(Debug, Clone, Copy, Dupe, PartialEq, Eq)]
328pub enum AssignOp {
329    Add,         // +=
330    Subtract,    // -=
331    Multiply,    // *=
332    Divide,      // /=
333    FloorDivide, // //=
334    Percent,     // %=
335    BitAnd,      // &=
336    BitOr,       // |=
337    BitXor,      // ^=
338    LeftShift,   // <<=
339    RightShift,  // >>=
340}
341
342#[derive(Debug, Copy, Clone, Dupe, Eq, PartialEq, Allocative)]
343pub enum Visibility {
344    Private,
345    Public,
346}
347
348#[derive(Debug, Clone)]
349pub struct DefP<P: AstPayload> {
350    pub name: AstAssignIdentP<P>,
351    pub params: Vec<AstParameterP<P>>,
352    pub return_type: Option<Box<AstTypeExprP<P>>>,
353    pub body: Box<AstStmtP<P>>,
354    pub payload: P::DefPayload,
355}
356
357impl<P: AstPayload> DefP<P> {
358    pub fn signature_span(&self) -> Span {
359        let mut span = self.name.span;
360        for param in &self.params {
361            span = span.merge(param.span);
362        }
363        if let Some(return_type) = &self.return_type {
364            span = span.merge(return_type.span);
365        }
366        span
367    }
368}
369
370#[derive(Debug, Clone)]
371pub struct ForP<P: AstPayload> {
372    pub var: AstAssignTargetP<P>,
373    pub over: AstExprP<P>,
374    pub body: Box<AstStmtP<P>>,
375}
376
377#[derive(Debug, Clone)]
378pub struct FStringP<P: AstPayload> {
379    /// A format string containing a `{}` marker for each expression to interpolate.
380    pub format: AstString,
381    /// The expressions to interpolate.
382    pub expressions: Vec<AstExprP<P>>,
383}
384
385/// A field in a struct definition.
386#[derive(Debug, Clone)]
387pub struct StructFieldP<P: AstPayload> {
388    /// Field name.
389    pub name: AstAssignIdentP<P>,
390    /// Field type annotation.
391    pub typ: AstTypeExprP<P>,
392    /// Default value (optional).
393    pub default: Option<AstExprP<P>>,
394}
395
396pub type AstStructFieldP<P> = Spanned<StructFieldP<P>>;
397pub type AstStructField = AstStructFieldP<AstNoPayload>;
398pub type StructField = StructFieldP<AstNoPayload>;
399
400/// A struct definition.
401#[derive(Debug, Clone)]
402pub struct StructP<P: AstPayload> {
403    /// Struct name.
404    pub name: AstAssignIdentP<P>,
405    /// Struct fields.
406    pub fields: Vec<AstStructFieldP<P>>,
407}
408
409/// A case arm in a match statement.
410#[derive(Debug, Clone)]
411pub struct CaseClauseP<P: AstPayload> {
412    /// The pattern to match against.
413    pub pattern: AstExprP<P>,
414    /// Optional guard condition.
415    pub guard: Option<AstExprP<P>>,
416    /// The body to execute if matched.
417    pub body: AstStmtP<P>,
418}
419
420pub type AstCaseClauseP<P> = Spanned<CaseClauseP<P>>;
421pub type AstCaseClause = AstCaseClauseP<AstNoPayload>;
422pub type CaseClause = CaseClauseP<AstNoPayload>;
423
424/// A match statement.
425#[derive(Debug, Clone)]
426pub struct MatchP<P: AstPayload> {
427    /// The expression to match against.
428    pub subject: AstExprP<P>,
429    /// The case clauses.
430    pub cases: Vec<AstCaseClauseP<P>>,
431}
432
433#[derive(Debug, Clone)]
434pub enum StmtP<P: AstPayload> {
435    Break,
436    Continue,
437    Pass,
438    Return(Option<AstExprP<P>>),
439    Yield(Option<AstExprP<P>>),
440    Expression(AstExprP<P>),
441    Assign(AssignP<P>),
442    AssignModify(AstAssignTargetP<P>, AssignOp, Box<AstExprP<P>>),
443    Statements(Vec<AstStmtP<P>>),
444    If(AstExprP<P>, Box<AstStmtP<P>>),
445    IfElse(AstExprP<P>, Box<(AstStmtP<P>, AstStmtP<P>)>),
446    For(ForP<P>),
447    Def(DefP<P>),
448    Load(LoadP<P>),
449    Struct(StructP<P>),
450    Match(MatchP<P>),
451}
452
453impl<P: AstPayload> ArgumentP<P> {
454    pub fn expr(&self) -> &AstExprP<P> {
455        match self {
456            ArgumentP::Positional(x) => x,
457            ArgumentP::Named(_, x) => x,
458            ArgumentP::Args(x) => x,
459            ArgumentP::KwArgs(x) => x,
460        }
461    }
462
463    pub fn expr_mut(&mut self) -> &mut AstExprP<P> {
464        match self {
465            ArgumentP::Positional(x) => x,
466            ArgumentP::Named(_, x) => x,
467            ArgumentP::Args(x) => x,
468            ArgumentP::KwArgs(x) => x,
469        }
470    }
471
472    /// Argument name if it is named.
473    pub fn name(&self) -> Option<&str> {
474        match self {
475            ArgumentP::Named(name, _) => Some(&name.node),
476            _ => None,
477        }
478    }
479}
480
481impl Display for BinOp {
482    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
483        match *self {
484            BinOp::Or => f.write_str(" or "),
485            BinOp::And => f.write_str(" and "),
486            BinOp::Equal => f.write_str(" == "),
487            BinOp::NotEqual => f.write_str(" != "),
488            BinOp::Less => f.write_str(" < "),
489            BinOp::Greater => f.write_str(" > "),
490            BinOp::LessOrEqual => f.write_str(" <= "),
491            BinOp::GreaterOrEqual => f.write_str(" >= "),
492            BinOp::In => f.write_str(" in "),
493            BinOp::NotIn => f.write_str(" not in "),
494            BinOp::Subtract => f.write_str(" - "),
495            BinOp::Add => f.write_str(" + "),
496            BinOp::Multiply => f.write_str(" * "),
497            BinOp::Percent => f.write_str(" % "),
498            BinOp::Divide => f.write_str(" / "),
499            BinOp::FloorDivide => f.write_str(" // "),
500            BinOp::BitAnd => f.write_str(" & "),
501            BinOp::BitOr => f.write_str(" | "),
502            BinOp::BitXor => f.write_str(" ^ "),
503            BinOp::LeftShift => f.write_str(" << "),
504            BinOp::RightShift => f.write_str(" >> "),
505        }
506    }
507}
508
509impl Display for AssignOp {
510    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
511        match *self {
512            AssignOp::Add => f.write_str(" += "),
513            AssignOp::Subtract => f.write_str(" -= "),
514            AssignOp::Multiply => f.write_str(" *= "),
515            AssignOp::Divide => f.write_str(" /= "),
516            AssignOp::FloorDivide => f.write_str(" //= "),
517            AssignOp::Percent => f.write_str(" %= "),
518            AssignOp::BitAnd => f.write_str(" &= "),
519            AssignOp::BitOr => f.write_str(" |= "),
520            AssignOp::BitXor => f.write_str(" ^= "),
521            AssignOp::LeftShift => f.write_str(" <<= "),
522            AssignOp::RightShift => f.write_str(" >>= "),
523        }
524    }
525}
526
527fn comma_separated_fmt<I, F>(
528    f: &mut Formatter<'_>,
529    v: &[I],
530    converter: F,
531    for_tuple: bool,
532) -> fmt::Result
533where
534    F: Fn(&I, &mut Formatter<'_>) -> fmt::Result,
535{
536    for (i, e) in v.iter().enumerate() {
537        f.write_str(if i == 0 { "" } else { ", " })?;
538        converter(e, f)?;
539    }
540    if v.len() == 1 && for_tuple {
541        f.write_str(",")?;
542    }
543    Ok(())
544}
545
546fn fmt_string_literal(f: &mut Formatter<'_>, s: &str) -> fmt::Result {
547    f.write_str("\"")?;
548    for c in s.chars() {
549        match c {
550            '\n' => f.write_str("\\n")?,
551            '\t' => f.write_str("\\t")?,
552            '\r' => f.write_str("\\r")?,
553            '\0' => f.write_str("\\0")?,
554            '"' => f.write_str("\\\"")?,
555            '\\' => f.write_str("\\\\")?,
556            x => f.write_str(&x.to_string())?,
557        }
558    }
559    f.write_str("\"")
560}
561
562impl Display for AstLiteral {
563    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
564        match self {
565            AstLiteral::Int(i) => write!(f, "{}", &i.node),
566            AstLiteral::Float(n) => write!(f, "{}", &n.node),
567            AstLiteral::String(s) => fmt_string_literal(f, &s.node),
568            AstLiteral::Ellipsis => f.write_str("..."),
569        }
570    }
571}
572
573impl Display for Expr {
574    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
575        match self {
576            Expr::Tuple(e) => {
577                f.write_str("(")?;
578                comma_separated_fmt(f, e, |x, f| write!(f, "{}", x.node), true)?;
579                f.write_str(")")
580            }
581            Expr::Dot(e, s) => write!(f, "{}.{}", e.node, s.node),
582            Expr::Lambda(LambdaP {
583                params,
584                body,
585                payload: _,
586            }) => {
587                f.write_str("(lambda ")?;
588                comma_separated_fmt(f, params, |x, f| write!(f, "{}", x.node), false)?;
589                f.write_str(": ")?;
590                write!(f, "{}", body.node)?;
591                f.write_str(")")
592            }
593            Expr::Call(e, args) => {
594                write!(f, "{}(", e.node)?;
595                for (i, x) in args.args.iter().enumerate() {
596                    if i != 0 {
597                        f.write_str(", ")?;
598                    }
599                    write!(f, "{}", x.node)?;
600                }
601                f.write_str(")")
602            }
603            Expr::Index(e_i) => {
604                let (e, i) = &**e_i;
605                write!(f, "{}[{}]", e.node, i.node)
606            }
607            Expr::Index2(a_i0_i1) => {
608                let (a, i0, i1) = &**a_i0_i1;
609                write!(f, "{}[{}, {}]", a.node, i0.node, i1.node)
610            }
611            Expr::Slice(e, i1, i2, i3) => {
612                write!(f, "{}[]", e.node)?;
613                if let Some(x) = i1 {
614                    write!(f, "{}:", x.node)?
615                } else {
616                    f.write_str(":")?
617                }
618                if let Some(x) = i2 {
619                    write!(f, "{}", x.node)?
620                }
621                if let Some(x) = i3 {
622                    write!(f, ":{}", x.node)?
623                }
624                Ok(())
625            }
626            Expr::Identifier(s) => Display::fmt(&s.node, f),
627            Expr::Not(e) => write!(f, "(not {})", e.node),
628            Expr::Minus(e) => write!(f, "-{}", e.node),
629            Expr::Plus(e) => write!(f, "+{}", e.node),
630            Expr::BitNot(e) => write!(f, "~{}", e.node),
631            Expr::Op(l, op, r) => write!(f, "({}{}{})", l.node, op, r.node),
632            Expr::If(cond_v1_v2) => {
633                let (cond, v1, v2) = &**cond_v1_v2;
634                write!(f, "({} if {} else {})", v1.node, cond.node, v2.node)
635            }
636            Expr::List(v) => {
637                f.write_str("[")?;
638                comma_separated_fmt(f, v, |x, f| write!(f, "{}", x.node), false)?;
639                f.write_str("]")
640            }
641            Expr::Dict(v) => {
642                f.write_str("{")?;
643                comma_separated_fmt(f, v, |x, f| write!(f, "{}: {}", x.0.node, x.1.node), false)?;
644                f.write_str("}")
645            }
646            Expr::Set(v) => {
647                f.write_str("{")?;
648                comma_separated_fmt(f, v, |x, f| write!(f, "{}", x.node), false)?;
649                f.write_str("}")
650            }
651            Expr::ListComprehension(e, for_, c) => {
652                write!(f, "[{}", e.node)?;
653                write!(f, "{for_}")?;
654                for x in c {
655                    write!(f, "{x}")?;
656                }
657                f.write_str("]")
658            }
659            Expr::SetComprehension(e, for_, c) => {
660                write!(f, "{{{}", e.node)?;
661                write!(f, "{for_}")?;
662                for x in c {
663                    write!(f, "{x}")?;
664                }
665                f.write_str("}")
666            }
667            Expr::DictComprehension(k_v, for_, c) => {
668                let (k, v) = &**k_v;
669                write!(f, "{{{}: {}", k.node, v.node)?;
670                write!(f, "{for_}")?;
671                for x in c {
672                    write!(f, "{x}")?;
673                }
674                f.write_str("}")
675            }
676            Expr::Literal(x) => write!(f, "{x}"),
677            Expr::FString(x) => {
678                // Write out the desugared form.
679                write!(f, "{}.format(", x.format.node)?;
680                comma_separated_fmt(f, &x.expressions, |x, f| write!(f, "{}", x.node), false)?;
681                f.write_str(")")
682            }
683        }
684    }
685}
686
687impl Display for TypeExpr {
688    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
689        Display::fmt(&self.expr.node, f)
690    }
691}
692
693impl Display for AssignTarget {
694    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
695        match self {
696            AssignTarget::Tuple(e) => {
697                f.write_str("(")?;
698                comma_separated_fmt(f, e, |x, f| write!(f, "{}", x.node), true)?;
699                f.write_str(")")
700            }
701            AssignTarget::Dot(e, s) => write!(f, "{}.{}", e.node, s.node),
702            AssignTarget::Index(e_i) => {
703                let (e, i) = &**e_i;
704                write!(f, "{}[{}]", e.node, i.node)
705            }
706            AssignTarget::Identifier(s) => write!(f, "{}", s.node),
707        }
708    }
709}
710
711impl Display for AssignIdent {
712    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
713        write!(f, "{}", self.ident)
714    }
715}
716
717impl Display for Ident {
718    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
719        write!(f, "{}", self.ident)
720    }
721}
722
723impl Display for Argument {
724    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
725        match self {
726            Argument::Positional(s) => write!(f, "{}", s.node),
727            Argument::Named(s, e) => write!(f, "{} = {}", s.node, e.node),
728            Argument::Args(s) => write!(f, "*{}", s.node),
729            Argument::KwArgs(s) => write!(f, "**{}", s.node),
730        }
731    }
732}
733
734impl Display for Parameter {
735    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
736        let (prefix, name, typ, default) = match self {
737            Parameter::Slash => return write!(f, "/"),
738            Parameter::Normal(s, t, e) => ("", s, t, e.as_ref()),
739            Parameter::NoArgs => return write!(f, "*"),
740            Parameter::Args(s, t) => ("*", s, t, None),
741            Parameter::KwArgs(s, t) => ("**", s, t, None),
742        };
743        write!(f, "{}{}", prefix, name.node)?;
744        if let Some(t) = typ {
745            write!(f, ": {}", t.node)?;
746        }
747        if let Some(d) = default {
748            write!(f, " = {}", d.node)?;
749        }
750        Ok(())
751    }
752}
753
754impl Display for ForClause {
755    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
756        write!(f, " for {} in {}", self.var.node, self.over.node)
757    }
758}
759
760impl Display for Clause {
761    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
762        match self {
763            Clause::For(x) => write!(f, "{x}"),
764            Clause::If(x) => write!(f, " if {}", x.node),
765        }
766    }
767}
768
769impl Stmt {
770    fn fmt_with_tab(&self, f: &mut Formatter<'_>, tab: String) -> fmt::Result {
771        match self {
772            Stmt::Break => writeln!(f, "{tab}break"),
773            Stmt::Continue => writeln!(f, "{tab}continue"),
774            Stmt::Pass => writeln!(f, "{tab}pass"),
775            Stmt::Return(Some(e)) => writeln!(f, "{}return {}", tab, e.node),
776            Stmt::Return(None) => writeln!(f, "{tab}return"),
777            Stmt::Yield(Some(e)) => writeln!(f, "{}yield {}", tab, e.node),
778            Stmt::Yield(None) => writeln!(f, "{tab}yield"),
779            Stmt::Expression(e) => writeln!(f, "{}{}", tab, e.node),
780            Stmt::Assign(AssignP { lhs, ty, rhs }) => {
781                write!(f, "{}{} ", tab, lhs.node)?;
782                if let Some(ty) = ty {
783                    write!(f, ": {} ", ty.node)?;
784                }
785                writeln!(f, "= {}", rhs.node)
786            }
787            Stmt::AssignModify(l, op, r) => writeln!(f, "{}{}{}{}", tab, l.node, op, r.node),
788            Stmt::Statements(v) => {
789                for s in v {
790                    s.node.fmt_with_tab(f, tab.clone())?;
791                }
792                Ok(())
793            }
794            Stmt::If(cond, suite) => {
795                writeln!(f, "{}if {}:", tab, cond.node)?;
796                suite.node.fmt_with_tab(f, tab + "  ")
797            }
798            Stmt::IfElse(cond, suite_1_2) => {
799                let (suite1, suite2) = &**suite_1_2;
800                writeln!(f, "{}if {}:", tab, cond.node)?;
801                suite1.node.fmt_with_tab(f, tab.clone() + "  ")?;
802                writeln!(f, "{tab}else:")?;
803                suite2.node.fmt_with_tab(f, tab + "  ")
804            }
805            Stmt::For(ForP { var, over, body }) => {
806                writeln!(f, "{}for {} in {}:", tab, var.node, over.node)?;
807                body.node.fmt_with_tab(f, tab + "  ")
808            }
809            Stmt::Def(DefP {
810                name,
811                params,
812                return_type,
813                body,
814                payload: _,
815            }) => {
816                write!(f, "{}def {}(", tab, name.node)?;
817                comma_separated_fmt(f, params, |x, f| write!(f, "{}", x.node), false)?;
818                f.write_str(")")?;
819                if let Some(rt) = return_type {
820                    write!(f, " -> {}", rt.node)?;
821                }
822                f.write_str(":\n")?;
823                body.node.fmt_with_tab(f, tab + "  ")
824            }
825            Stmt::Load(load) => {
826                write!(f, "{tab}load(")?;
827                fmt_string_literal(f, &load.module.node)?;
828                f.write_str(", ")?;
829                comma_separated_fmt(
830                    f,
831                    &load.args,
832                    |x, f| {
833                        write!(f, "{} = ", x.local.node)?;
834                        fmt_string_literal(f, &(x.their.node))
835                    },
836                    false,
837                )?;
838                f.write_str(")\n")
839            }
840            Stmt::Struct(StructP { name, fields }) => {
841                writeln!(f, "{}struct {}:", tab, name.node)?;
842                for field in fields {
843                    write!(f, "{}  {}: {}", tab, field.node.name.node, field.node.typ.expr.node)?;
844                    if let Some(default) = &field.node.default {
845                        write!(f, " = {}", default.node)?;
846                    }
847                    writeln!(f)?;
848                }
849                Ok(())
850            }
851            Stmt::Match(MatchP { subject, cases }) => {
852                writeln!(f, "{}match {}:", tab, subject.node)?;
853                for case in cases {
854                    write!(f, "{}  case {}", tab, case.node.pattern.node)?;
855                    if let Some(guard) = &case.node.guard {
856                        write!(f, " if {}", guard.node)?;
857                    }
858                    writeln!(f, ":")?;
859                    case.node.body.node.fmt_with_tab(f, tab.clone() + "    ")?;
860                }
861                Ok(())
862            }
863        }
864    }
865}
866
867impl Display for Stmt {
868    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
869        self.fmt_with_tab(f, "".to_owned())
870    }
871}
blueprint_starlark_syntax/syntax/ast.rs

blueprint_starlark_syntax/syntax/
ast.rs
