1use miette::{SourceOffset, SourceSpan};
2
3#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] pub struct Node {
5 pub start: usize,
7
8 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#[derive(Debug, Clone)]
35pub struct Module {
36 pub node: Node,
37 pub body: Vec<Statement>,
38}
39
40#[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#[derive(Debug, Clone)]
167pub struct Raise {
168 pub node: Node,
169 pub exc: Option<Expression>,
170 pub cause: Option<Expression>,
171}
172
173#[derive(Debug, Clone)]
175pub struct Break {
176 pub node: Node,
177}
178
179#[derive(Debug, Clone)]
181pub struct Continue {
182 pub node: Node,
183}
184
185#[derive(Debug, Clone)]
187pub struct Import {
188 pub node: Node,
189 pub names: Vec<Alias>,
190}
191
192#[derive(Debug, Clone)]
194pub struct Alias {
195 pub node: Node,
196 pub name: String,
197 pub asname: Option<String>,
198}
199
200#[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#[derive(Debug, Clone)]
211pub struct Global {
212 pub node: Node,
213 pub names: Vec<String>,
214}
215
216#[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#[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 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#[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#[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#[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#[derive(Debug, Clone)]
419pub struct NamedExpression {
420 pub node: Node,
421 pub target: Box<Expression>,
422 pub value: Box<Expression>,
423}
424
425#[derive(Debug, Clone)]
427pub struct Yield {
428 pub node: Node,
429 pub value: Option<Box<Expression>>,
430}
431
432#[derive(Debug, Clone)]
434pub struct YieldFrom {
435 pub node: Node,
436 pub value: Box<Expression>,
437}
438
439#[derive(Debug, Clone)]
441pub struct Starred {
442 pub node: Node,
443 pub value: Box<Expression>,
444}
445
446#[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#[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#[derive(Debug, Clone)]
488pub struct Attribute {
489 pub node: Node,
490 pub value: Box<Expression>,
491 pub attr: String,
492}
493
494#[derive(Debug, Clone)]
496pub struct Subscript {
497 pub node: Node,
498 pub value: Box<Expression>,
499 pub slice: Box<Expression>,
500}
501
502#[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#[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#[derive(Debug, Clone)]
532pub struct Await {
533 pub node: Node,
534 pub value: Box<Expression>,
535}
536
537#[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#[derive(Debug, Clone)]
562pub struct Lambda {
563 pub node: Node,
564 pub args: Arguments,
565 pub body: Box<Expression>,
566}
567
568#[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#[derive(Debug, Clone)]
583pub struct Arg {
584 pub node: Node,
585 pub arg: String,
586 pub annotation: Option<Expression>,
587}
588
589#[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#[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#[derive(Debug, Clone)]
609pub struct JoinedStr {
610 pub node: Node,
611 pub values: Vec<Expression>,
612}
613
614#[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#[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#[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#[derive(Debug, Clone)]
650pub struct With {
651 pub node: Node,
652 pub items: Vec<WithItem>,
653 pub body: Vec<Statement>,
654}
655
656#[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#[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#[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#[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#[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#[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#[derive(Debug, Clone)]
719pub struct Match {
720 pub node: Node,
721 pub subject: Box<Expression>,
722 pub cases: Vec<MatchCase>,
723}
724
725#[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}