Skip to main content

oak_python/ast/
mod.rs

1use core::range::Range;
2use serde::{Deserialize, Serialize};
3
4/// Python 源文件的根节点
5#[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/// Python 程序
13#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
14pub struct Program {
15    pub statements: Vec<Statement>,
16}
17
18/// 语句
19#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
20pub enum Statement {
21    /// 函数定义
22    FunctionDef { name: String, parameters: Vec<Parameter>, return_type: Option<Type>, body: Vec<Statement> },
23    /// 类定义
24    ClassDef { name: String, bases: Vec<Expression>, body: Vec<Statement> },
25    /// 变量赋值
26    Assignment { target: Expression, value: Expression },
27    /// 复合赋值 (+=, -=, 等)
28    AugmentedAssignment { target: Expression, operator: AugmentedOperator, value: Expression },
29    /// 表达式语句
30    Expression(Expression),
31    /// 返回语句
32    Return(Option<Expression>),
33    /// 条件语句
34    If { test: Expression, body: Vec<Statement>, orelse: Vec<Statement> },
35    /// for 循环
36    For { target: Expression, iter: Expression, body: Vec<Statement>, orelse: Vec<Statement> },
37    /// while 循环
38    While { test: Expression, body: Vec<Statement>, orelse: Vec<Statement> },
39    /// break 语句
40    Break,
41    /// continue 语句
42    Continue,
43    /// pass 语句
44    Pass,
45    /// import 语句
46    Import { names: Vec<ImportName> },
47    /// from import 语句
48    ImportFrom { module: Option<String>, names: Vec<ImportName> },
49    /// try 语句
50    Try { body: Vec<Statement>, handlers: Vec<ExceptHandler>, orelse: Vec<Statement>, finalbody: Vec<Statement> },
51    /// raise 语句
52    Raise { exc: Option<Expression>, cause: Option<Expression> },
53    /// with 语句
54    With { items: Vec<WithItem>, body: Vec<Statement> },
55}
56
57/// 表达式
58#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
59pub enum Expression {
60    /// 字面量
61    Literal(Literal),
62    /// 标识符
63    Name(String),
64    /// 二元运算
65    BinaryOp { left: Box<Expression>, operator: BinaryOperator, right: Box<Expression> },
66    /// 一元运算
67    UnaryOp { operator: UnaryOperator, operand: Box<Expression> },
68    /// 布尔运算 (and, or)
69    BoolOp { operator: BoolOperator, values: Vec<Expression> },
70    /// 比较运算
71    Compare { left: Box<Expression>, ops: Vec<CompareOperator>, comparators: Vec<Expression> },
72    /// 函数调用
73    Call { func: Box<Expression>, args: Vec<Expression>, keywords: Vec<Keyword> },
74    /// 属性访问
75    Attribute { value: Box<Expression>, attr: String },
76    /// 下标访问
77    Subscript { value: Box<Expression>, slice: Box<Expression> },
78    /// 列表
79    List { elts: Vec<Expression> },
80    /// 元组
81    Tuple { elts: Vec<Expression> },
82    /// 字典
83    Dict { keys: Vec<Option<Expression>>, values: Vec<Expression> },
84    /// 集合
85    Set { elts: Vec<Expression> },
86    /// 列表推导式
87    ListComp { elt: Box<Expression>, generators: Vec<Comprehension> },
88    /// 字典推导式
89    DictComp { key: Box<Expression>, value: Box<Expression>, generators: Vec<Comprehension> },
90    /// 集合推导式
91    SetComp { elt: Box<Expression>, generators: Vec<Comprehension> },
92    /// lambda 表达式
93    Lambda { args: Vec<Parameter>, body: Box<Expression> },
94    /// 条件表达式 (三元运算符)
95    IfExp { test: Box<Expression>, body: Box<Expression>, orelse: Box<Expression> },
96}
97
98/// 字面量
99#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
100pub enum Literal {
101    /// 整数
102    Integer(i64),
103    /// 浮点数
104    Float(f64),
105    /// 字符串
106    String(String),
107    /// 布尔值
108    Boolean(bool),
109    /// None
110    None,
111}
112
113/// 二元运算符
114#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
115pub enum BinaryOperator {
116    Add,      // +
117    Sub,      // -
118    Mult,     // *
119    Div,      // /
120    FloorDiv, // //
121    Mod,      // %
122    Pow,      // **
123    LShift,   // <<
124    RShift,   // >>
125    BitOr,    // |
126    BitXor,   // ^
127    BitAnd,   // &
128}
129
130/// 一元运算符
131#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
132pub enum UnaryOperator {
133    Invert, // ~
134    Not,    // not
135    UAdd,   // +
136    USub,   // -
137}
138
139/// 布尔运算符
140#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
141pub enum BoolOperator {
142    And, // and
143    Or,  // or
144}
145
146/// 比较运算符
147#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
148pub enum CompareOperator {
149    Eq,    // ==
150    NotEq, // !=
151    Lt,    // <
152    LtE,   // <=
153    Gt,    // >
154    GtE,   // >=
155    Is,    // is
156    IsNot, // is not
157    In,    // in
158    NotIn, // not in
159}
160
161/// 复合赋值运算符
162#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
163pub enum AugmentedOperator {
164    Add,      // +=
165    Sub,      // -=
166    Mult,     // *=
167    Div,      // /=
168    FloorDiv, // //=
169    Mod,      // %=
170    Pow,      // **=
171    LShift,   // <<=
172    RShift,   // >>=
173    BitOr,    // |=
174    BitXor,   // ^=
175    BitAnd,   // &=
176}
177
178/// 函数参数
179#[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/// 类型注解
187#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
188pub enum Type {
189    /// 基本类型名称
190    Name(String),
191    /// 泛型类型
192    Generic { name: String, args: Vec<Type> },
193    /// 联合类型
194    Union(Vec<Type>),
195    /// 可选类型
196    Optional(Box<Type>),
197}
198
199/// 关键字参数
200#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
201pub struct Keyword {
202    pub arg: Option<String>,
203    pub value: Expression,
204}
205
206/// 推导式
207#[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/// import 名称
216#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
217pub struct ImportName {
218    pub name: String,
219    pub asname: Option<String>,
220}
221
222/// 异常处理器
223#[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/// with 语句项
231#[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    /// 创建一个新的程序
239    pub fn new() -> Self {
240        Self { statements: Vec::new() }
241    }
242
243    /// 添加语句
244    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    /// 创建一个标识符表达式
257    pub fn name(name: impl Into<String>) -> Self {
258        Self::Name(name.into())
259    }
260
261    /// 创建一个字符串字面量表达式
262    pub fn string(value: impl Into<String>) -> Self {
263        Self::Literal(Literal::String(value.into()))
264    }
265
266    /// 创建一个整数字面量表达式
267    pub fn integer(value: i64) -> Self {
268        Self::Literal(Literal::Integer(value))
269    }
270
271    /// 创建一个浮点数字面量表达式
272    pub fn float(value: f64) -> Self {
273        Self::Literal(Literal::Float(value))
274    }
275
276    /// 创建一个布尔字面量表达式
277    pub fn boolean(value: bool) -> Self {
278        Self::Literal(Literal::Boolean(value))
279    }
280
281    /// 创建一个 None 字面量表达式
282    pub fn none() -> Self {
283        Self::Literal(Literal::None)
284    }
285}
286
287impl Statement {
288    /// 创建一个函数定义语句
289    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    /// 创建一个赋值语句
294    pub fn assignment(target: Expression, value: Expression) -> Self {
295        Self::Assignment { target, value }
296    }
297
298    /// 创建一个表达式语句
299    pub fn expression(expr: Expression) -> Self {
300        Self::Expression(expr)
301    }
302
303    /// 创建一个返回语句
304    pub fn return_stmt(value: Option<Expression>) -> Self {
305        Self::Return(value)
306    }
307}