enderpy_python_parser/parser/
ast.rs

1use miette::{SourceOffset, SourceSpan};
2
3#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] // #[serde(tag = "type")]
4pub struct Node {
5    /// Start offset in source
6    pub start: usize,
7
8    /// End offset in source
9    pub end: usize,
10}
11
12impl Node {
13    pub fn new(start: usize, end: usize) -> Self {
14        Self { start, end }
15    }
16
17    pub fn len(&self) -> usize {
18        self.end - self.start
19    }
20}
21
22pub trait GetNode {
23    fn get_node(&self) -> Node;
24}
25
26impl From<Node> for SourceSpan {
27    fn from(val: Node) -> Self {
28        Self::new(SourceOffset::from(val.start), SourceOffset::from(val.len()))
29    }
30}
31
32// The following structs are used to represent the AST
33// https://docs.python.org/3/library/ast.html#abstract-grammar
34#[derive(Debug, Clone)]
35pub struct Module {
36    pub node: Node,
37    pub body: Vec<Statement>,
38}
39
40// Use box to reduce the enum size
41#[derive(Debug, Clone)]
42pub enum Statement {
43    AssignStatement(Assign),
44    AnnAssignStatement(AnnAssign),
45    AugAssignStatement(AugAssign),
46    ExpressionStatement(Expression),
47    Assert(Assert),
48    Pass(Pass),
49    Delete(Delete),
50    Return(Return),
51    Raise(Raise),
52    Break(Break),
53    Continue(Continue),
54    Import(Import),
55    ImportFrom(ImportFrom),
56    Global(Global),
57    Nonlocal(Nonlocal),
58    IfStatement(If),
59    WhileStatement(While),
60    ForStatement(For),
61    WithStatement(With),
62    TryStatement(Try),
63    TryStarStatement(TryStar),
64    FunctionDef(FunctionDef),
65    ClassDef(ClassDef),
66    Match(Match),
67}
68
69impl GetNode for Statement {
70    fn get_node(&self) -> Node {
71        match self {
72            Statement::AssignStatement(s) => s.node,
73            Statement::AnnAssignStatement(s) => s.node,
74            Statement::AugAssignStatement(s) => s.node,
75            Statement::ExpressionStatement(s) => s.get_node(),
76            Statement::Assert(s) => s.node,
77            Statement::Pass(s) => s.node,
78            Statement::Delete(s) => s.node,
79            Statement::Return(s) => s.node,
80            Statement::Raise(s) => s.node,
81            Statement::Break(s) => s.node,
82            Statement::Continue(s) => s.node,
83            Statement::Import(s) => s.node,
84            Statement::ImportFrom(s) => s.node,
85            Statement::Global(s) => s.node,
86            Statement::Nonlocal(s) => s.node,
87            Statement::IfStatement(s) => s.node,
88            Statement::WhileStatement(s) => s.node,
89            Statement::ForStatement(s) => s.node,
90            Statement::WithStatement(s) => s.node,
91            Statement::TryStatement(s) => s.node,
92            Statement::TryStarStatement(s) => s.node,
93            Statement::FunctionDef(s) => s.node,
94            Statement::ClassDef(s) => s.node,
95            Statement::Match(s) => s.node,
96        }
97    }
98}
99
100#[derive(Debug, Clone)]
101pub struct Assign {
102    pub node: Node,
103    pub targets: Vec<Expression>,
104    pub value: Expression,
105}
106
107#[derive(Debug, Clone)]
108pub struct AnnAssign {
109    pub node: Node,
110    pub target: Expression,
111    pub annotation: Expression,
112    pub value: Option<Expression>,
113    pub simple: bool,
114}
115
116#[derive(Debug, Clone)]
117pub struct AugAssign {
118    pub node: Node,
119    pub target: Expression,
120    pub op: AugAssignOp,
121    pub value: Expression,
122}
123
124#[derive(Debug, Clone)]
125pub enum AugAssignOp {
126    Add,
127    Sub,
128    Mult,
129    MatMult,
130    Div,
131    Mod,
132    Pow,
133    LShift,
134    RShift,
135    BitOr,
136    BitXor,
137    BitAnd,
138    FloorDiv,
139}
140
141#[derive(Debug, Clone)]
142pub struct Assert {
143    pub node: Node,
144    pub test: Expression,
145    pub msg: Option<Expression>,
146}
147
148#[derive(Debug, Clone)]
149pub struct Pass {
150    pub node: Node,
151}
152
153#[derive(Debug, Clone)]
154pub struct Delete {
155    pub node: Node,
156    pub targets: Vec<Expression>,
157}
158
159#[derive(Debug, Clone)]
160pub struct Return {
161    pub node: Node,
162    pub value: Option<Expression>,
163}
164
165// https://docs.python.org/3/library/ast.html#ast.Raise
166#[derive(Debug, Clone)]
167pub struct Raise {
168    pub node: Node,
169    pub exc: Option<Expression>,
170    pub cause: Option<Expression>,
171}
172
173// https://docs.python.org/3/library/ast.html#ast.Break
174#[derive(Debug, Clone)]
175pub struct Break {
176    pub node: Node,
177}
178
179// https://docs.python.org/3/library/ast.html#ast.Continue
180#[derive(Debug, Clone)]
181pub struct Continue {
182    pub node: Node,
183}
184
185// https://docs.python.org/3/library/ast.html#ast.Import
186#[derive(Debug, Clone)]
187pub struct Import {
188    pub node: Node,
189    pub names: Vec<Alias>,
190}
191
192// https://docs.python.org/3/library/ast.html#ast.alias
193#[derive(Debug, Clone)]
194pub struct Alias {
195    pub node: Node,
196    pub name: String,
197    pub asname: Option<String>,
198}
199
200// https://docs.python.org/3/library/ast.html#ast.ImportFrom
201#[derive(Debug, Clone)]
202pub struct ImportFrom {
203    pub node: Node,
204    pub module: String,
205    pub names: Vec<Alias>,
206    pub level: usize,
207}
208
209// https://docs.python.org/3/library/ast.html#ast.Global
210#[derive(Debug, Clone)]
211pub struct Global {
212    pub node: Node,
213    pub names: Vec<String>,
214}
215
216// https://docs.python.org/3/library/ast.html#ast.Nonlocal
217#[derive(Debug, Clone)]
218pub struct Nonlocal {
219    pub node: Node,
220    pub names: Vec<String>,
221}
222
223#[derive(Debug, Clone)]
224pub enum Expression {
225    Constant(Box<Constant>),
226    List(Box<List>),
227    Tuple(Box<Tuple>),
228    Dict(Box<Dict>),
229    Set(Box<Set>),
230    Name(Box<Name>),
231    BoolOp(Box<BoolOperation>),
232    UnaryOp(Box<UnaryOperation>),
233    BinOp(Box<BinOp>),
234    NamedExpr(Box<NamedExpression>),
235    Yield(Box<Yield>),
236    YieldFrom(Box<YieldFrom>),
237    Starred(Box<Starred>),
238    Generator(Box<Generator>),
239    ListComp(Box<ListComp>),
240    SetComp(Box<SetComp>),
241    DictComp(Box<DictComp>),
242    Attribute(Box<Attribute>),
243    Subscript(Box<Subscript>),
244    Slice(Box<Slice>),
245    Call(Box<Call>),
246    Await(Box<Await>),
247    Compare(Box<Compare>),
248    Lambda(Box<Lambda>),
249    IfExp(Box<IfExp>),
250    JoinedStr(Box<JoinedStr>),
251    FormattedValue(Box<FormattedValue>),
252}
253
254impl GetNode for Expression {
255    fn get_node(&self) -> Node {
256        match self {
257            Expression::Constant(c) => c.node,
258            Expression::List(l) => l.node,
259            Expression::Tuple(t) => t.node,
260            Expression::Dict(d) => d.node,
261            Expression::Set(s) => s.node,
262            Expression::Name(n) => n.node,
263            Expression::BoolOp(b) => b.node,
264            Expression::UnaryOp(u) => u.node,
265            Expression::BinOp(b) => b.node,
266            Expression::NamedExpr(n) => n.node,
267            Expression::Yield(y) => y.node,
268            Expression::YieldFrom(y) => y.node,
269            Expression::Starred(s) => s.node,
270            Expression::Generator(g) => g.node,
271            Expression::ListComp(l) => l.node,
272            Expression::SetComp(s) => s.node,
273            Expression::DictComp(d) => d.node,
274            Expression::Attribute(a) => a.node,
275            Expression::Subscript(s) => s.node,
276            Expression::Slice(s) => s.node,
277            Expression::Call(c) => c.node,
278            Expression::Await(a) => a.node,
279            Expression::Compare(c) => c.node,
280            Expression::Lambda(l) => l.node,
281            Expression::IfExp(i) => i.node,
282            Expression::JoinedStr(j) => j.node,
283            Expression::FormattedValue(f) => f.node,
284        }
285    }
286}
287
288// https://docs.python.org/3/reference/expressions.html#atom-identifiers
289#[derive(Debug, Clone)]
290pub struct Name {
291    pub node: Node,
292    pub id: String,
293}
294
295#[derive(Clone, Debug, PartialEq)]
296pub struct Constant {
297    pub node: Node,
298    pub value: ConstantValue,
299}
300
301#[derive(Clone, Debug, PartialEq)]
302pub enum ConstantValue {
303    None,
304    Bool(bool),
305    Str(String),
306    Bytes(Vec<u8>),
307    Tuple(Vec<Constant>),
308    // Numbers are string because we don't care about the value rn.
309    Int(String),
310    Float(String),
311    Complex { real: String, imaginary: String },
312}
313
314#[derive(Debug, Clone)]
315pub struct List {
316    pub node: Node,
317    pub elements: Vec<Expression>,
318}
319
320#[derive(Debug, Clone)]
321pub struct Tuple {
322    pub node: Node,
323    pub elements: Vec<Expression>,
324}
325
326#[derive(Debug, Clone)]
327pub struct Dict {
328    pub node: Node,
329    pub keys: Vec<Expression>,
330    pub values: Vec<Expression>,
331}
332
333#[derive(Debug, Clone)]
334pub struct Set {
335    pub node: Node,
336    pub elements: Vec<Expression>,
337}
338
339// https://docs.python.org/3/library/ast.html#ast.BoolOp
340#[derive(Debug, Clone)]
341pub struct BoolOperation {
342    pub node: Node,
343    pub op: BooleanOperator,
344    pub values: Vec<Expression>,
345}
346
347#[derive(Debug, Clone)]
348pub enum BooleanOperator {
349    And,
350    Or,
351}
352
353// https://docs.python.org/3/library/ast.html#ast.UnaryOp
354#[derive(Debug, Clone)]
355pub struct UnaryOperation {
356    pub node: Node,
357    pub op: UnaryOperator,
358    pub operand: Box<Expression>,
359}
360
361#[derive(Debug, Clone)]
362pub enum UnaryOperator {
363    Not,
364    Invert,
365    UAdd,
366    USub,
367}
368
369// https://docs.python.org/3/library/ast.html#ast.BinOp
370#[derive(Debug, Clone)]
371pub struct BinOp {
372    pub node: Node,
373    pub op: BinaryOperator,
374    pub left: Box<Expression>,
375    pub right: Box<Expression>,
376}
377
378#[derive(Debug, Clone)]
379pub enum BinaryOperator {
380    Add,
381    Sub,
382    Mult,
383    MatMult,
384    Div,
385    Mod,
386    Pow,
387    LShift,
388    RShift,
389    BitOr,
390    BitXor,
391    BitAnd,
392    FloorDiv,
393}
394
395impl std::fmt::Display for BinaryOperator {
396    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
397        let op_str = match self {
398            BinaryOperator::Add => "+",
399            BinaryOperator::Sub => "-",
400            BinaryOperator::Mult => "*",
401            BinaryOperator::MatMult => "@",
402            BinaryOperator::Div => "/",
403            BinaryOperator::Mod => "%",
404            BinaryOperator::Pow => "**",
405            BinaryOperator::LShift => "<<",
406            BinaryOperator::RShift => ">>",
407            BinaryOperator::BitOr => "|",
408            BinaryOperator::BitXor => "^",
409            BinaryOperator::BitAnd => "&",
410            BinaryOperator::FloorDiv => "//",
411        };
412
413        write!(f, "{}", op_str)
414    }
415}
416
417// https://docs.python.org/3/library/ast.html#ast.NamedExpr
418#[derive(Debug, Clone)]
419pub struct NamedExpression {
420    pub node: Node,
421    pub target: Box<Expression>,
422    pub value: Box<Expression>,
423}
424
425// https://docs.python.org/3/library/ast.html#ast.Yield
426#[derive(Debug, Clone)]
427pub struct Yield {
428    pub node: Node,
429    pub value: Option<Box<Expression>>,
430}
431
432// https://docs.python.org/3/library/ast.html#ast.YieldFrom
433#[derive(Debug, Clone)]
434pub struct YieldFrom {
435    pub node: Node,
436    pub value: Box<Expression>,
437}
438
439// https://docs.python.org/3/library/ast.html#ast.Starred
440#[derive(Debug, Clone)]
441pub struct Starred {
442    pub node: Node,
443    pub value: Box<Expression>,
444}
445
446// https://docs.python.org/3/library/ast.html#ast.GeneratorExp
447#[derive(Debug, Clone)]
448pub struct Generator {
449    pub node: Node,
450    pub element: Box<Expression>,
451    pub generators: Vec<Comprehension>,
452}
453
454#[derive(Debug, Clone)]
455pub struct ListComp {
456    pub node: Node,
457    pub element: Box<Expression>,
458    pub generators: Vec<Comprehension>,
459}
460
461#[derive(Debug, Clone)]
462pub struct SetComp {
463    pub node: Node,
464    pub element: Box<Expression>,
465    pub generators: Vec<Comprehension>,
466}
467
468#[derive(Debug, Clone)]
469pub struct DictComp {
470    pub node: Node,
471    pub key: Box<Expression>,
472    pub value: Box<Expression>,
473    pub generators: Vec<Comprehension>,
474}
475
476// https://docs.python.org/3/library/ast.html#ast.comprehension
477#[derive(Debug, Clone)]
478pub struct Comprehension {
479    pub node: Node,
480    pub target: Box<Expression>,
481    pub iter: Box<Expression>,
482    pub ifs: Vec<Expression>,
483    pub is_async: bool,
484}
485
486// https://docs.python.org/3/library/ast.html#ast.Attribute
487#[derive(Debug, Clone)]
488pub struct Attribute {
489    pub node: Node,
490    pub value: Box<Expression>,
491    pub attr: String,
492}
493
494// https://docs.python.org/3/library/ast.html#ast.Subscript
495#[derive(Debug, Clone)]
496pub struct Subscript {
497    pub node: Node,
498    pub value: Box<Expression>,
499    pub slice: Box<Expression>,
500}
501
502// https://docs.python.org/3/library/ast.html#ast.Slice
503// can be used for Subscript
504#[derive(Debug, Clone)]
505pub struct Slice {
506    pub node: Node,
507    pub lower: Option<Box<Expression>>,
508    pub upper: Option<Box<Expression>>,
509    pub step: Option<Box<Expression>>,
510}
511
512// https://docs.python.org/3/library/ast.html#ast.Call
513#[derive(Debug, Clone)]
514pub struct Call {
515    pub node: Node,
516    pub func: Box<Expression>,
517    pub args: Vec<Expression>,
518    pub keywords: Vec<Keyword>,
519    pub starargs: Option<Box<Expression>>,
520    pub kwargs: Option<Box<Expression>>,
521}
522
523#[derive(Debug, Clone)]
524pub struct Keyword {
525    pub node: Node,
526    pub arg: Option<String>,
527    pub value: Box<Expression>,
528}
529
530// https://docs.python.org/3/library/ast.html#ast.Await
531#[derive(Debug, Clone)]
532pub struct Await {
533    pub node: Node,
534    pub value: Box<Expression>,
535}
536
537// https://docs.python.org/3/library/ast.html#ast.Compare
538#[derive(Debug, Clone)]
539pub struct Compare {
540    pub node: Node,
541    pub left: Box<Expression>,
542    pub ops: Vec<ComparisonOperator>,
543    pub comparators: Vec<Expression>,
544}
545
546#[derive(Debug, Clone)]
547pub enum ComparisonOperator {
548    Eq,
549    NotEq,
550    Lt,
551    LtE,
552    Gt,
553    GtE,
554    Is,
555    IsNot,
556    In,
557    NotIn,
558}
559
560// https://docs.python.org/3/library/ast.html#ast.Lambda
561#[derive(Debug, Clone)]
562pub struct Lambda {
563    pub node: Node,
564    pub args: Arguments,
565    pub body: Box<Expression>,
566}
567
568// https://docs.python.org/3/library/ast.html#ast.arguments
569#[derive(Debug, Clone)]
570pub struct Arguments {
571    pub node: Node,
572    pub posonlyargs: Vec<Arg>,
573    pub args: Vec<Arg>,
574    pub vararg: Option<Arg>,
575    pub kwonlyargs: Vec<Arg>,
576    pub kw_defaults: Vec<Option<Expression>>,
577    pub kwarg: Option<Arg>,
578    pub defaults: Vec<Expression>,
579}
580
581// https://docs.python.org/3/library/ast.html#ast.arg
582#[derive(Debug, Clone)]
583pub struct Arg {
584    pub node: Node,
585    pub arg: String,
586    pub annotation: Option<Expression>,
587}
588
589// https://docs.python.org/3/library/ast.html#ast.IfExp
590#[derive(Debug, Clone)]
591pub struct IfExp {
592    pub node: Node,
593    pub test: Box<Expression>,
594    pub body: Box<Expression>,
595    pub orelse: Box<Expression>,
596}
597
598// https://docs.python.org/3/library/ast.html#ast.FormattedValue
599#[derive(Debug, Clone)]
600pub struct FormattedValue {
601    pub node: Node,
602    pub value: Box<Expression>,
603    pub conversion: Option<i32>,
604    pub format_spec: Option<Box<Expression>>,
605}
606
607// https://docs.python.org/3/library/ast.html#ast.JoinedStr
608#[derive(Debug, Clone)]
609pub struct JoinedStr {
610    pub node: Node,
611    pub values: Vec<Expression>,
612}
613
614// https://docs.python.org/3/library/ast.html#ast.If
615#[derive(Debug, Clone)]
616pub struct If {
617    pub node: Node,
618    pub test: Box<Expression>,
619    pub body: Vec<Statement>,
620    pub orelse: Vec<Statement>,
621}
622
623impl If {
624    pub fn update_orelse(&mut self, other_or_else: Vec<Statement>) {
625        self.orelse = other_or_else;
626    }
627}
628
629// https://docs.python.org/3/library/ast.html#ast.While
630#[derive(Debug, Clone)]
631pub struct While {
632    pub node: Node,
633    pub test: Box<Expression>,
634    pub body: Vec<Statement>,
635    pub orelse: Vec<Statement>,
636}
637
638// https://docs.python.org/3/library/ast.html#ast.For
639#[derive(Debug, Clone)]
640pub struct For {
641    pub node: Node,
642    pub target: Box<Expression>,
643    pub iter: Box<Expression>,
644    pub body: Vec<Statement>,
645    pub orelse: Vec<Statement>,
646}
647
648// https://docs.python.org/3/library/ast.html#ast.With
649#[derive(Debug, Clone)]
650pub struct With {
651    pub node: Node,
652    pub items: Vec<WithItem>,
653    pub body: Vec<Statement>,
654}
655
656// https://docs.python.org/3/library/ast.html#ast.withitem
657// can be used for With
658#[derive(Debug, Clone)]
659pub struct WithItem {
660    pub node: Node,
661    pub context_expr: Box<Expression>,
662    pub optional_vars: Option<Box<Expression>>,
663}
664
665// https://docs.python.org/3/library/ast.html#ast.Try
666#[derive(Debug, Clone)]
667pub struct Try {
668    pub node: Node,
669    pub body: Vec<Statement>,
670    pub handlers: Vec<ExceptHandler>,
671    pub orelse: Vec<Statement>,
672    pub finalbody: Vec<Statement>,
673}
674
675// https://docs.python.org/3/library/ast.html#ast.TryStar
676#[derive(Debug, Clone)]
677pub struct TryStar {
678    pub node: Node,
679    pub body: Vec<Statement>,
680    pub handlers: Vec<ExceptHandler>,
681    pub orelse: Vec<Statement>,
682    pub finalbody: Vec<Statement>,
683}
684
685// https://docs.python.org/3/library/ast.html#ast.ExceptHandler
686#[derive(Debug, Clone)]
687pub struct ExceptHandler {
688    pub node: Node,
689    pub typ: Option<Box<Expression>>,
690    pub name: Option<String>,
691    pub body: Vec<Statement>,
692}
693
694// https://docs.python.org/3/library/ast.html#functiondef
695#[derive(Debug, Clone)]
696pub struct FunctionDef {
697    pub node: Node,
698    pub name: String,
699    pub args: Arguments,
700    pub body: Vec<Statement>,
701    pub decorator_list: Vec<Expression>,
702    pub returns: Option<Box<Expression>>,
703    pub type_comment: Option<String>,
704}
705
706// https://docs.python.org/3/library/ast.html#ast.ClassDef
707#[derive(Debug, Clone)]
708pub struct ClassDef {
709    pub node: Node,
710    pub name: String,
711    pub bases: Vec<Expression>,
712    pub keywords: Vec<Keyword>,
713    pub body: Vec<Statement>,
714    pub decorator_list: Vec<Expression>,
715}
716
717// https://docs.python.org/3/library/ast.html#ast.Match
718#[derive(Debug, Clone)]
719pub struct Match {
720    pub node: Node,
721    pub subject: Box<Expression>,
722    pub cases: Vec<MatchCase>,
723}
724
725// https://docs.python.org/3/library/ast.html#ast.match_case
726#[derive(Debug, Clone)]
727pub struct MatchCase {
728    pub node: Node,
729    pub pattern: Box<MatchPattern>,
730    pub guard: Option<Box<Expression>>,
731    pub body: Vec<Statement>,
732}
733
734#[derive(Debug, Clone)]
735pub enum MatchPattern {
736    MatchValue(MatchValue),
737    MatchSingleton(Box<Expression>),
738    MatchSequence(Vec<MatchPattern>),
739    MatchStar(Box<Expression>),
740    MatchMapping(MatchMapping),
741    MatchAs(MatchAs),
742    MatchClass(MatchClass),
743    MatchOr(Vec<MatchPattern>),
744}
745
746#[derive(Debug, Clone)]
747pub struct MatchValue {
748    pub node: Node,
749    pub value: Box<Expression>,
750}
751
752#[derive(Debug, Clone)]
753pub struct MatchAs {
754    pub node: Node,
755    pub name: Option<String>,
756    pub pattern: Option<Box<MatchPattern>>,
757}
758
759#[derive(Debug, Clone)]
760pub struct MatchMapping {
761    pub node: Node,
762    pub keys: Vec<Expression>,
763    pub patterns: Vec<MatchPattern>,
764    pub rest: Option<String>,
765}
766
767#[derive(Debug, Clone)]
768pub struct MatchClass {
769    pub node: Node,
770    pub cls: Box<Expression>,
771    pub patterns: Vec<MatchPattern>,
772    pub kwd_attrs: Vec<String>,
773    pub kwd_patterns: Vec<MatchPattern>,
774}