1use crate::{metadata::Metadata, types::Type, NodeId, Visibility};
7use serde::{Deserialize, Serialize};
8
9#[allow(clippy::module_name_repetitions)]
11#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
12pub enum PythonHIR {
13 Module {
15 name: String,
17 body: Vec<PythonHIR>,
19 meta: Metadata,
21 },
22
23 Function {
25 id: NodeId,
27 name: String,
29 params: Vec<Parameter>,
31 return_type: Option<Type>,
33 body: Vec<PythonHIR>,
35 decorators: Vec<String>,
37 visibility: Visibility,
39 meta: Metadata,
41 },
42
43 Class {
45 id: NodeId,
47 name: String,
49 bases: Vec<String>,
51 body: Vec<PythonHIR>,
53 decorators: Vec<String>,
55 meta: Metadata,
57 },
58
59 Call {
61 id: NodeId,
63 callee: Box<PythonHIR>,
65 args: Vec<PythonHIR>,
67 kwargs: Vec<(String, PythonHIR)>,
69 inferred_type: Option<Type>,
71 meta: Metadata,
73 },
74
75 Variable {
77 id: NodeId,
79 name: String,
81 inferred_type: Option<Type>,
83 meta: Metadata,
85 },
86
87 Assign {
89 id: NodeId,
91 target: String,
93 value: Box<PythonHIR>,
95 type_annotation: Option<Type>,
97 meta: Metadata,
99 },
100
101 Return {
103 id: NodeId,
105 value: Option<Box<PythonHIR>>,
107 meta: Metadata,
109 },
110
111 If {
113 id: NodeId,
115 condition: Box<PythonHIR>,
117 then_branch: Vec<PythonHIR>,
119 else_branch: Vec<PythonHIR>,
121 meta: Metadata,
123 },
124
125 For {
127 id: NodeId,
129 target: String,
131 iter: Box<PythonHIR>,
133 body: Vec<PythonHIR>,
135 orelse: Vec<PythonHIR>,
137 meta: Metadata,
139 },
140
141 While {
143 id: NodeId,
145 condition: Box<PythonHIR>,
147 body: Vec<PythonHIR>,
149 orelse: Vec<PythonHIR>,
151 meta: Metadata,
153 },
154
155 BinOp {
157 id: NodeId,
159 op: BinOp,
161 left: Box<PythonHIR>,
163 right: Box<PythonHIR>,
165 inferred_type: Option<Type>,
167 meta: Metadata,
169 },
170
171 UnaryOp {
173 id: NodeId,
175 op: UnaryOp,
177 operand: Box<PythonHIR>,
179 inferred_type: Option<Type>,
181 meta: Metadata,
183 },
184
185 Literal {
187 id: NodeId,
189 value: Literal,
191 meta: Metadata,
193 },
194
195 ListComp {
197 id: NodeId,
199 element: Box<PythonHIR>,
201 generators: Vec<Comprehension>,
203 meta: Metadata,
205 },
206
207 Attribute {
209 id: NodeId,
211 object: Box<PythonHIR>,
213 attr: String,
215 inferred_type: Option<Type>,
217 meta: Metadata,
219 },
220
221 Subscript {
223 id: NodeId,
225 object: Box<PythonHIR>,
227 index: Box<PythonHIR>,
229 inferred_type: Option<Type>,
231 meta: Metadata,
233 },
234}
235
236#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
238pub struct Parameter {
239 pub name: String,
241 pub type_annotation: Option<Type>,
243 pub default: Option<String>,
245}
246
247#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
249pub enum BinOp {
250 Add,
252 Sub,
254 Mul,
256 Div,
258 FloorDiv,
260 Mod,
262 Pow,
264 Eq,
266 NotEq,
268 Lt,
270 Le,
272 Gt,
274 Ge,
276 And,
278 Or,
280}
281
282#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
284pub enum UnaryOp {
285 Not,
287 Neg,
289 Pos,
291}
292
293#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
295pub enum Literal {
296 Int(i64),
298 Float(f64),
300 Str(String),
302 Bool(bool),
304 None,
306}
307
308#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
310pub struct Comprehension {
311 pub target: String,
313 pub iter: Box<PythonHIR>,
315 pub ifs: Vec<PythonHIR>,
317}
318
319impl PythonHIR {
320 #[must_use]
322 pub const fn id(&self) -> Option<NodeId> {
323 match self {
324 Self::Module { .. } => None,
325 Self::Function { id, .. }
326 | Self::Class { id, .. }
327 | Self::Call { id, .. }
328 | Self::Variable { id, .. }
329 | Self::Assign { id, .. }
330 | Self::Return { id, .. }
331 | Self::If { id, .. }
332 | Self::For { id, .. }
333 | Self::While { id, .. }
334 | Self::BinOp { id, .. }
335 | Self::UnaryOp { id, .. }
336 | Self::Literal { id, .. }
337 | Self::ListComp { id, .. }
338 | Self::Attribute { id, .. }
339 | Self::Subscript { id, .. } => Some(*id),
340 }
341 }
342
343 #[must_use]
345 pub const fn metadata(&self) -> &Metadata {
346 match self {
347 Self::Module { meta, .. }
348 | Self::Function { meta, .. }
349 | Self::Class { meta, .. }
350 | Self::Call { meta, .. }
351 | Self::Variable { meta, .. }
352 | Self::Assign { meta, .. }
353 | Self::Return { meta, .. }
354 | Self::If { meta, .. }
355 | Self::For { meta, .. }
356 | Self::While { meta, .. }
357 | Self::BinOp { meta, .. }
358 | Self::UnaryOp { meta, .. }
359 | Self::Literal { meta, .. }
360 | Self::ListComp { meta, .. }
361 | Self::Attribute { meta, .. }
362 | Self::Subscript { meta, .. } => meta,
363 }
364 }
365}
366
367#[cfg(test)]
368mod tests {
369 use super::*;
370
371 #[test]
372 fn test_python_function_creation() {
373 let func = PythonHIR::Function {
374 id: NodeId::new(1),
375 name: "test_func".to_owned(),
376 params: vec![],
377 return_type: None,
378 body: vec![],
379 decorators: vec![],
380 visibility: Visibility::Public,
381 meta: Metadata::new(),
382 };
383
384 assert_eq!(func.id(), Some(NodeId::new(1)));
385 }
386
387 #[test]
388 fn test_python_call_creation() {
389 let call = PythonHIR::Call {
390 id: NodeId::new(2),
391 callee: Box::new(PythonHIR::Variable {
392 id: NodeId::new(3),
393 name: "len".to_owned(),
394 inferred_type: None,
395 meta: Metadata::new(),
396 }),
397 args: vec![],
398 kwargs: vec![],
399 inferred_type: None,
400 meta: Metadata::new(),
401 };
402
403 assert_eq!(call.id(), Some(NodeId::new(2)));
404 }
405}