1use core::range::Range;
2use serde::{Deserialize, Serialize};
3
4#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
6pub struct PythonRoot {
7 pub program: Program,
8 #[serde(with = "oak_core::serde_range")]
9 pub span: Range<usize>,
10}
11
12#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
14pub struct Program {
15 pub statements: Vec<Statement>,
16}
17
18#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
20pub enum Statement {
21 FunctionDef { name: String, parameters: Vec<Parameter>, return_type: Option<Type>, body: Vec<Statement> },
23 ClassDef { name: String, bases: Vec<Expression>, body: Vec<Statement> },
25 Assignment { target: Expression, value: Expression },
27 AugmentedAssignment { target: Expression, operator: AugmentedOperator, value: Expression },
29 Expression(Expression),
31 Return(Option<Expression>),
33 If { test: Expression, body: Vec<Statement>, orelse: Vec<Statement> },
35 For { target: Expression, iter: Expression, body: Vec<Statement>, orelse: Vec<Statement> },
37 While { test: Expression, body: Vec<Statement>, orelse: Vec<Statement> },
39 Break,
41 Continue,
43 Pass,
45 Import { names: Vec<ImportName> },
47 ImportFrom { module: Option<String>, names: Vec<ImportName> },
49 Try { body: Vec<Statement>, handlers: Vec<ExceptHandler>, orelse: Vec<Statement>, finalbody: Vec<Statement> },
51 Raise { exc: Option<Expression>, cause: Option<Expression> },
53 With { items: Vec<WithItem>, body: Vec<Statement> },
55}
56
57#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
59pub enum Expression {
60 Literal(Literal),
62 Name(String),
64 BinaryOp { left: Box<Expression>, operator: BinaryOperator, right: Box<Expression> },
66 UnaryOp { operator: UnaryOperator, operand: Box<Expression> },
68 BoolOp { operator: BoolOperator, values: Vec<Expression> },
70 Compare { left: Box<Expression>, ops: Vec<CompareOperator>, comparators: Vec<Expression> },
72 Call { func: Box<Expression>, args: Vec<Expression>, keywords: Vec<Keyword> },
74 Attribute { value: Box<Expression>, attr: String },
76 Subscript { value: Box<Expression>, slice: Box<Expression> },
78 List { elts: Vec<Expression> },
80 Tuple { elts: Vec<Expression> },
82 Dict { keys: Vec<Option<Expression>>, values: Vec<Expression> },
84 Set { elts: Vec<Expression> },
86 ListComp { elt: Box<Expression>, generators: Vec<Comprehension> },
88 DictComp { key: Box<Expression>, value: Box<Expression>, generators: Vec<Comprehension> },
90 SetComp { elt: Box<Expression>, generators: Vec<Comprehension> },
92 Lambda { args: Vec<Parameter>, body: Box<Expression> },
94 IfExp { test: Box<Expression>, body: Box<Expression>, orelse: Box<Expression> },
96}
97
98#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
100pub enum Literal {
101 Integer(i64),
103 Float(f64),
105 String(String),
107 Boolean(bool),
109 None,
111}
112
113#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
115pub enum BinaryOperator {
116 Add, Sub, Mult, Div, FloorDiv, Mod, Pow, LShift, RShift, BitOr, BitXor, BitAnd, }
129
130#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
132pub enum UnaryOperator {
133 Invert, Not, UAdd, USub, }
138
139#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
141pub enum BoolOperator {
142 And, Or, }
145
146#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
148pub enum CompareOperator {
149 Eq, NotEq, Lt, LtE, Gt, GtE, Is, IsNot, In, NotIn, }
160
161#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
163pub enum AugmentedOperator {
164 Add, Sub, Mult, Div, FloorDiv, Mod, Pow, LShift, RShift, BitOr, BitXor, BitAnd, }
177
178#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
180pub struct Parameter {
181 pub name: String,
182 pub annotation: Option<Type>,
183 pub default: Option<Expression>,
184}
185
186#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
188pub enum Type {
189 Name(String),
191 Generic { name: String, args: Vec<Type> },
193 Union(Vec<Type>),
195 Optional(Box<Type>),
197}
198
199#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
201pub struct Keyword {
202 pub arg: Option<String>,
203 pub value: Expression,
204}
205
206#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
208pub struct Comprehension {
209 pub target: Expression,
210 pub iter: Expression,
211 pub ifs: Vec<Expression>,
212 pub is_async: bool,
213}
214
215#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
217pub struct ImportName {
218 pub name: String,
219 pub asname: Option<String>,
220}
221
222#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
224pub struct ExceptHandler {
225 pub type_: Option<Expression>,
226 pub name: Option<String>,
227 pub body: Vec<Statement>,
228}
229
230#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
232pub struct WithItem {
233 pub context_expr: Expression,
234 pub optional_vars: Option<Expression>,
235}
236
237impl Program {
238 pub fn new() -> Self {
240 Self { statements: Vec::new() }
241 }
242
243 pub fn add_statement(&mut self, statement: Statement) {
245 self.statements.push(statement);
246 }
247}
248
249impl Default for Program {
250 fn default() -> Self {
251 Self::new()
252 }
253}
254
255impl Expression {
256 pub fn name(name: impl Into<String>) -> Self {
258 Self::Name(name.into())
259 }
260
261 pub fn string(value: impl Into<String>) -> Self {
263 Self::Literal(Literal::String(value.into()))
264 }
265
266 pub fn integer(value: i64) -> Self {
268 Self::Literal(Literal::Integer(value))
269 }
270
271 pub fn float(value: f64) -> Self {
273 Self::Literal(Literal::Float(value))
274 }
275
276 pub fn boolean(value: bool) -> Self {
278 Self::Literal(Literal::Boolean(value))
279 }
280
281 pub fn none() -> Self {
283 Self::Literal(Literal::None)
284 }
285}
286
287impl Statement {
288 pub fn function_def(name: impl Into<String>, parameters: Vec<Parameter>, return_type: Option<Type>, body: Vec<Statement>) -> Self {
290 Self::FunctionDef { name: name.into(), parameters, return_type, body }
291 }
292
293 pub fn assignment(target: Expression, value: Expression) -> Self {
295 Self::Assignment { target, value }
296 }
297
298 pub fn expression(expr: Expression) -> Self {
300 Self::Expression(expr)
301 }
302
303 pub fn return_stmt(value: Option<Expression>) -> Self {
305 Self::Return(value)
306 }
307}