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

blueprint_starlark_syntax/syntax/
ast.rs
