oak_python/builder/
mod.rs

1use crate::{
2    ast::{BinaryOperator, ExceptHandler, Expression, ImportName, Keyword, Literal, Parameter, Program, PythonRoot, Statement, WithItem},
3    language::PythonLanguage,
4    lexer::token_type::PythonTokenType,
5    parser::{PythonParser, element_type::PythonElementType},
6};
7use oak_core::{Builder, BuilderCache, GreenNode, GreenTree, OakDiagnostics, OakError, Parser, SourceText, TextEdit, builder::BuildOutput, source::Source};
8
9/// AST builder for the Python language.
10pub struct PythonBuilder<'config> {
11    /// Reference to the language configuration
12    config: &'config PythonLanguage,
13}
14
15impl<'config> PythonBuilder<'config> {
16    /// Parses an f-string's text content into a list of expressions.
17    fn parse_fstring_text(&self, text: &str) -> Result<Vec<Expression>, OakError> {
18        let mut values = Vec::new();
19        let mut i = 0;
20        let bytes = text.as_bytes();
21
22        // Skip prefix (f, rf, fr)
23        while i < bytes.len() && (bytes[i] == b'f' || bytes[i] == b'F' || bytes[i] == b'r' || bytes[i] == b'R') {
24            i += 1;
25        }
26
27        // Skip opening quotes
28        if i >= bytes.len() {
29            return Ok(values);
30        }
31        let quote_char = bytes[i];
32        let mut quote_count = 0;
33        while i < bytes.len() && bytes[i] == quote_char && quote_count < 3 {
34            i += 1;
35            quote_count += 1;
36        }
37
38        let end_index = if bytes.len() >= quote_count { bytes.len() - quote_count } else { bytes.len() };
39        let mut start = i;
40
41        while i < end_index {
42            if bytes[i] == b'{' {
43                // Potential expression
44                if i + 1 < end_index && bytes[i + 1] == b'{' {
45                    // Escaped {{
46                    i += 2;
47                    continue;
48                }
49
50                // Add literal part before {
51                if i > start {
52                    let s = &text[start..i];
53                    values.push(Expression::Literal(Literal::String(s.to_string())));
54                }
55
56                // Find matching }
57                i += 1; // skip {
58                let expr_start = i;
59                let mut brace_level = 1;
60                while i < end_index && brace_level > 0 {
61                    if bytes[i] == b'{' {
62                        brace_level += 1
63                    }
64                    else if bytes[i] == b'}' {
65                        brace_level -= 1
66                    }
67                    i += 1
68                }
69
70                if brace_level == 0 {
71                    let expr_text = &text[expr_start..i - 1];
72                    let expr = self.parse_single_expression(expr_text)?;
73                    values.push(Expression::FormattedValue { value: Box::new(expr), conversion: 0, format_spec: None });
74                    start = i
75                }
76                else {
77                    // Unmatched {
78                    start = i
79                }
80            }
81            else if bytes[i] == b'}' {
82                if i + 1 < end_index && bytes[i + 1] == b'}' {
83                    // Escaped }}
84                    i += 2;
85                    continue;
86                }
87                i += 1
88            }
89            else {
90                i += 1
91            }
92        }
93
94        if start < end_index {
95            let s = &text[start..end_index];
96            values.push(Expression::Literal(Literal::String(s.to_string())))
97        }
98
99        Ok(values)
100    }
101
102    /// Parses a single expression from text.
103    fn parse_single_expression(&self, text: &str) -> Result<Expression, OakError> {
104        let parser = PythonParser::new(self.config);
105        let mut cache = oak_core::parser::session::ParseSession::<PythonLanguage>::default();
106        let source = SourceText::new(text.to_string());
107        let output = parser.parse(&source, &[], &mut cache);
108
109        match output.result {
110            Ok(green) => {
111                for child in green.children() {
112                    if let GreenTree::Node(node) = child {
113                        if node.kind == PythonElementType::Expr {
114                            for subchild in node.children() {
115                                if let GreenTree::Node(expr_node) = subchild {
116                                    if !expr_node.kind.is_trivia() {
117                                        return self.build_expression(expr_node, 0, &source);
118                                    }
119                                }
120                            }
121                        }
122                    }
123                }
124                Ok(Expression::Name(text.to_string()))
125            }
126            Err(e) => Err(e),
127        }
128    }
129
130    /// Creates a new Python builder.
131    pub fn new(config: &'config PythonLanguage) -> Self {
132        Self { config }
133    }
134}
135
136impl<'config> Builder<PythonLanguage> for PythonBuilder<'config> {
137    fn build<'a, S: Source + ?Sized>(&self, source: &S, edits: &[TextEdit], _cache: &'a mut impl BuilderCache<PythonLanguage>) -> BuildOutput<PythonLanguage> {
138        let parser = PythonParser::new(self.config);
139
140        let mut parse_cache = oak_core::parser::session::ParseSession::<PythonLanguage>::default();
141        let parse_result = parser.parse(source, edits, &mut parse_cache);
142
143        match parse_result.result {
144            Ok(green_tree) => {
145                let source_text = SourceText::new(source.get_text_in((0..source.length()).into()).into_owned());
146                match self.build_root(green_tree, &source_text) {
147                    Ok(ast_root) => OakDiagnostics { result: Ok(ast_root), diagnostics: parse_result.diagnostics },
148                    Err(build_error) => {
149                        let mut diagnostics = parse_result.diagnostics;
150                        diagnostics.push(build_error.clone());
151                        OakDiagnostics { result: Err(build_error), diagnostics }
152                    }
153                }
154            }
155            Err(e) => OakDiagnostics { result: Err(e), diagnostics: parse_result.diagnostics },
156        }
157    }
158}
159
160impl<'config> PythonBuilder<'config> {
161    /// Builds a Python root from a green tree.
162    pub fn build_root(&self, green_tree: &GreenNode<PythonLanguage>, source: &SourceText) -> Result<PythonRoot, OakError> {
163        let mut statements = Vec::new();
164        let mut current_offset = 0;
165        let mut pending_decorators = Vec::new();
166
167        for child in green_tree.children() {
168            let child_len = child.len() as usize;
169            match child {
170                GreenTree::Node(node) => {
171                    if !node.kind.is_trivia() {
172                        if node.kind == PythonElementType::Decorator {
173                            // Collect decorator expression
174                            let mut decorator_offset = current_offset;
175                            for d_child in node.children() {
176                                if let GreenTree::Node(expr_node) = d_child {
177                                    if !expr_node.kind.is_trivia() {
178                                        pending_decorators.push(self.build_expression(expr_node, decorator_offset, source)?)
179                                    }
180                                }
181                                decorator_offset += d_child.len() as usize
182                            }
183                        }
184                        else {
185                            if let Some(mut stmt) = self.build_statement(node, current_offset, source)? {
186                                // Attach decorators to function or class def
187                                match &mut stmt {
188                                    Statement::FunctionDef { decorators, .. } | Statement::AsyncFunctionDef { decorators, .. } | Statement::ClassDef { decorators, .. } => *decorators = std::mem::take(&mut pending_decorators),
189                                    _ => {
190                                        // If it's not a function or class def, decorators are invalid here
191                                        // but for now we just drop them or we could issue a warning
192                                        pending_decorators.clear()
193                                    }
194                                }
195                                statements.push(stmt)
196                            }
197                        }
198                    }
199                }
200                _ => {}
201            }
202            current_offset += child_len
203        }
204
205        Ok(PythonRoot { program: Program { statements }, span: (0..green_tree.text_len() as usize).into() })
206    }
207
208    /// Builds a statement from a green node.
209    fn build_statement(&self, node: &GreenNode<PythonLanguage>, offset: usize, source: &SourceText) -> Result<Option<Statement>, OakError> {
210        match node.kind {
211            PythonElementType::FunctionDef | PythonElementType::AsyncFunctionDef => {
212                let is_async = node.kind == PythonElementType::AsyncFunctionDef;
213                let mut name = String::new();
214                let mut parameters = Vec::new();
215                let mut body = Vec::new();
216                let mut current_offset = offset;
217
218                for child in node.children() {
219                    let child_len = child.len() as usize;
220                    match child {
221                        GreenTree::Leaf(leaf) if leaf.kind == PythonTokenType::Identifier => name = source.get_text_in((current_offset..current_offset + leaf.length as usize).into()).trim().to_string(),
222                        GreenTree::Node(n) if n.kind == PythonElementType::Arguments => parameters = self.build_parameters(n, current_offset, source)?,
223                        GreenTree::Node(n) if n.kind == PythonElementType::Suite => body = self.build_suite(n, current_offset, source)?,
224                        _ => {}
225                    }
226                    current_offset += child_len
227                }
228                if is_async { Ok(Some(Statement::AsyncFunctionDef { decorators: Vec::new(), name, parameters, return_type: None, body })) } else { Ok(Some(Statement::FunctionDef { decorators: Vec::new(), name, parameters, return_type: None, body })) }
229            }
230            PythonElementType::ClassDef => {
231                let mut name = String::new();
232                let mut bases = Vec::new();
233                let mut body = Vec::new();
234                let mut current_offset = offset;
235
236                for child in node.children() {
237                    let child_len = child.len() as usize;
238                    match child {
239                        GreenTree::Leaf(leaf) if leaf.kind == PythonTokenType::Identifier => name = source.get_text_in((current_offset..current_offset + leaf.length as usize).into()).trim().to_string(),
240                        GreenTree::Node(n) if n.kind == PythonElementType::Tuple || n.kind == PythonElementType::List => {
241                            // Simple bases handling: if it's a tuple or list in class def, treat as bases
242                            let expr = self.build_expression(n, current_offset, source)?;
243                            match expr {
244                                Expression::Tuple { elts } => bases = elts,
245                                Expression::List { elts } => bases = elts,
246                                _ => bases.push(expr),
247                            }
248                        }
249                        GreenTree::Node(n) if n.kind == PythonElementType::Suite => body = self.build_suite(n, current_offset, source)?,
250                        GreenTree::Node(n) => {
251                            // Try to see if it's an expression (base class)
252                            if !n.kind.is_trivia() && n.kind != PythonElementType::Suite {
253                                let expr = self.build_expression(n, current_offset, source)?;
254                                bases.push(expr)
255                            }
256                        }
257                        _ => {}
258                    }
259                    current_offset += child_len
260                }
261                Ok(Some(Statement::ClassDef { decorators: Vec::new(), name, bases, body }))
262            }
263            PythonElementType::Return => {
264                let mut value = None;
265                let mut current_offset = offset;
266                for child in node.children() {
267                    let child_len = child.len() as usize;
268                    if let GreenTree::Node(n) = child {
269                        if !n.kind.is_trivia() {
270                            value = Some(self.build_expression(n, current_offset, source)?)
271                        }
272                    }
273                    current_offset += child_len
274                }
275                Ok(Some(Statement::Return(value)))
276            }
277            PythonElementType::AssignStmt => {
278                let mut exprs = Vec::new();
279                let mut current_offset = offset;
280
281                for child in node.children() {
282                    let child_len = child.len() as usize;
283                    match child {
284                        GreenTree::Node(n) => {
285                            if !n.kind.is_trivia() {
286                                exprs.push(self.build_expression(n, current_offset, source)?)
287                            }
288                        }
289                        _ => {}
290                    }
291                    current_offset += child_len
292                }
293
294                if exprs.len() >= 2 {
295                    // Python supports multiple assignments: a = b = c = 0
296                    // In our AST, we'll just handle the first two for now to match the Statement::Assignment definition
297                    // or we could change the AST to support multiple targets.
298                    // Given the current AST: Assignment { target: Expression, value: Expression }
299                    // We'll treat it as target = value.
300                    let target = exprs[0].clone();
301                    let value = exprs[exprs.len() - 1].clone();
302                    Ok(Some(Statement::Assignment { target, value }))
303                }
304                else {
305                    Ok(None)
306                }
307            }
308
309            PythonElementType::Expr => {
310                let mut current_offset = offset;
311                for child in node.children() {
312                    let child_len = child.len() as usize;
313                    match child {
314                        GreenTree::Node(n) => {
315                            if !n.kind.is_trivia() {
316                                if n.kind == PythonElementType::AssignStmt {
317                                    return self.build_statement(n, current_offset, source);
318                                }
319                                return Ok(Some(Statement::Expression(self.build_expression(n, current_offset, source)?)));
320                            }
321                        }
322                        _ => {}
323                    }
324                    current_offset += child_len
325                }
326                Ok(None)
327            }
328            PythonElementType::If => {
329                let mut test = None;
330                let mut body = Vec::new();
331                let mut orelse = Vec::new();
332                let mut current_offset = offset;
333
334                for child in node.children() {
335                    let child_len = child.len() as usize;
336                    if let GreenTree::Node(n) = child {
337                        if !n.kind.is_trivia() {
338                            if test.is_none() && n.kind != PythonElementType::Suite {
339                                test = Some(self.build_expression(n, current_offset, source)?)
340                            }
341                            else if body.is_empty() && n.kind == PythonElementType::Suite {
342                                body = self.build_suite(n, current_offset, source)?
343                            }
344                            else if n.kind == PythonElementType::Suite {
345                                orelse = self.build_suite(n, current_offset, source)?
346                            }
347                        }
348                    }
349                    current_offset += child_len
350                }
351
352                Ok(Some(Statement::If { test: test.unwrap_or(Expression::Literal(Literal::Boolean(true))), body, orelse }))
353            }
354            PythonElementType::While => {
355                let mut test = None;
356                let mut body = Vec::new();
357                let mut current_offset = offset;
358
359                for child in node.children() {
360                    let child_len = child.len() as usize;
361                    if let GreenTree::Node(n) = child {
362                        if !n.kind.is_trivia() {
363                            if test.is_none() && n.kind != PythonElementType::Suite {
364                                test = Some(self.build_expression(n, current_offset, source)?)
365                            }
366                            else if n.kind == PythonElementType::Suite {
367                                body = self.build_suite(n, current_offset, source)?
368                            }
369                        }
370                    }
371                    current_offset += child_len
372                }
373
374                Ok(Some(Statement::While { test: test.unwrap_or(Expression::Literal(Literal::Boolean(true))), body, orelse: Vec::new() }))
375            }
376            PythonElementType::For | PythonElementType::AsyncFor => {
377                let is_async = node.kind == PythonElementType::AsyncFor;
378                let mut target = None;
379                let mut iter = None;
380                let mut body = Vec::new();
381                let mut orelse = Vec::new();
382                let mut current_offset = offset;
383
384                for child in node.children() {
385                    let child_len = child.len() as usize;
386                    if let GreenTree::Node(n) = child {
387                        if !n.kind.is_trivia() {
388                            if target.is_none() && n.kind != PythonElementType::Suite && n.kind != PythonElementType::InKeyword.into() {
389                                target = Some(self.build_expression(n, current_offset, source)?)
390                            }
391                            else if iter.is_none() && n.kind != PythonElementType::Suite && n.kind != PythonElementType::InKeyword.into() {
392                                iter = Some(self.build_expression(n, current_offset, source)?)
393                            }
394                            else if n.kind == PythonElementType::Suite {
395                                if body.is_empty() { body = self.build_suite(n, current_offset, source)? } else { orelse = self.build_suite(n, current_offset, source)? }
396                            }
397                        }
398                    }
399                    current_offset += child_len
400                }
401
402                let target = target.unwrap_or(Expression::Name("invalid_target".to_string()));
403                let iter = iter.unwrap_or(Expression::Name("invalid_iter".to_string()));
404                if is_async { Ok(Some(Statement::AsyncFor { target, iter, body, orelse })) } else { Ok(Some(Statement::For { target, iter, body, orelse })) }
405            }
406            PythonElementType::Pass => Ok(Some(Statement::Pass)),
407            PythonElementType::Break => Ok(Some(Statement::Break)),
408            PythonElementType::Continue => Ok(Some(Statement::Continue)),
409            PythonElementType::Raise => {
410                let mut exc = None;
411                let mut cause = None;
412                let mut current_offset = offset;
413                for child in node.children() {
414                    let child_len = child.len() as usize;
415                    if let GreenTree::Node(n) = child {
416                        if !n.kind.is_trivia() {
417                            let expr = self.build_expression(n, current_offset, source)?;
418                            if exc.is_none() { exc = Some(expr) } else { cause = Some(expr) }
419                        }
420                    }
421                    current_offset += child_len
422                }
423                Ok(Some(Statement::Raise { exc, cause }))
424            }
425            PythonElementType::Assert => {
426                let mut test = None;
427                let mut msg = None;
428                let mut current_offset = offset;
429                for child in node.children() {
430                    let child_len = child.len() as usize;
431                    if let GreenTree::Node(n) = child {
432                        if !n.kind.is_trivia() {
433                            let expr = self.build_expression(n, current_offset, source)?;
434                            if test.is_none() { test = Some(expr) } else { msg = Some(expr) }
435                        }
436                    }
437                    current_offset += child_len
438                }
439                Ok(Some(Statement::Assert { test: test.unwrap_or(Expression::Literal(crate::ast::Literal::Boolean(true))), msg }))
440            }
441            PythonElementType::Import => {
442                let mut names = Vec::new();
443                let mut current_offset = offset;
444                for child in node.children() {
445                    let child_len = child.len() as usize;
446                    if let GreenTree::Node(n) = child {
447                        if n.kind == PythonElementType::Alias {
448                            names.push(self.build_import_name(n, current_offset, source)?)
449                        }
450                    }
451                    current_offset += child_len
452                }
453                Ok(Some(Statement::Import { names }))
454            }
455            PythonElementType::ImportFrom => {
456                let mut module = None;
457                let mut names = Vec::new();
458                let mut current_offset = offset;
459                for child in node.children() {
460                    let child_len = child.len() as usize;
461                    match child {
462                        GreenTree::Leaf(leaf) if leaf.kind == PythonTokenType::Identifier => {
463                            if module.is_none() {
464                                module = Some(source.get_text_in((current_offset..current_offset + leaf.length as usize).into()).trim().to_string())
465                            }
466                        }
467                        GreenTree::Node(n) if n.kind == PythonElementType::Alias => names.push(self.build_import_name(n, current_offset, source)?),
468                        _ => {}
469                    }
470                    current_offset += child_len
471                }
472                Ok(Some(Statement::ImportFrom { module, names }))
473            }
474            PythonElementType::Global => {
475                let mut names = Vec::new();
476                let mut current_offset = offset;
477                for child in node.children() {
478                    let child_len = child.len() as usize;
479                    if let GreenTree::Leaf(leaf) = child {
480                        if leaf.kind == PythonTokenType::Identifier {
481                            names.push(source.get_text_in((current_offset..current_offset + leaf.length as usize).into()).trim().to_string())
482                        }
483                    }
484                    current_offset += child_len
485                }
486                Ok(Some(Statement::Global { names }))
487            }
488            PythonElementType::Nonlocal => {
489                let mut names = Vec::new();
490                let mut current_offset = offset;
491                for child in node.children() {
492                    let child_len = child.len() as usize;
493                    if let GreenTree::Leaf(leaf) = child {
494                        if leaf.kind == PythonTokenType::Identifier {
495                            names.push(source.get_text_in((current_offset..current_offset + leaf.length as usize).into()).trim().to_string())
496                        }
497                    }
498                    current_offset += child_len
499                }
500                Ok(Some(Statement::Nonlocal { names }))
501            }
502            PythonElementType::Try => {
503                let mut body = Vec::new();
504                let mut handlers = Vec::new();
505                let mut orelse = Vec::new();
506                let mut finalbody = Vec::new();
507                let mut current_offset = offset;
508
509                for child in node.children() {
510                    let child_len = child.len() as usize;
511                    if let GreenTree::Node(n) = child {
512                        if !n.kind.is_trivia() {
513                            match n.kind {
514                                PythonElementType::Suite if body.is_empty() && handlers.is_empty() => body = self.build_suite(n, current_offset, source)?,
515                                PythonElementType::ExceptHandler => handlers.push(self.build_except_handler(n, current_offset, source)?),
516                                PythonElementType::Suite if !handlers.is_empty() && orelse.is_empty() => orelse = self.build_suite(n, current_offset, source)?,
517                                PythonElementType::Suite if finalbody.is_empty() => finalbody = self.build_suite(n, current_offset, source)?,
518                                _ => {}
519                            }
520                        }
521                    }
522                    current_offset += child_len
523                }
524                Ok(Some(Statement::Try { body, handlers, orelse, finalbody }))
525            }
526            PythonElementType::With => {
527                let mut items = Vec::new();
528                let mut body = Vec::new();
529                let mut current_offset = offset;
530
531                for child in node.children() {
532                    let child_len = child.len() as usize;
533                    if let GreenTree::Node(n) = child {
534                        if !n.kind.is_trivia() {
535                            if n.kind == PythonElementType::WithItem {
536                                items.push(self.build_with_item(n, current_offset, source)?)
537                            }
538                            else if n.kind == PythonElementType::Suite {
539                                body = self.build_suite(n, current_offset, source)?
540                            }
541                        }
542                    }
543                    current_offset += child_len
544                }
545                Ok(Some(Statement::With { items, body }))
546            }
547            PythonElementType::Suite => {
548                // Suites are handled by build_suite
549                Ok(None)
550            }
551            _ => Ok(None),
552        }
553    }
554
555    /// Builds a suite (list of statements) from a green node.
556    fn build_suite(&self, node: &GreenNode<PythonLanguage>, offset: usize, source: &SourceText) -> Result<Vec<Statement>, OakError> {
557        let mut statements = Vec::new();
558        let mut current_offset = offset;
559
560        for child in node.children() {
561            let child_len = child.len() as usize;
562            match child {
563                GreenTree::Node(n) => {
564                    if !n.kind.is_trivia() {
565                        if let Some(stmt) = self.build_statement(n, current_offset, source)? {
566                            statements.push(stmt)
567                        }
568                    }
569                }
570                _ => {}
571            }
572            current_offset += child_len
573        }
574        Ok(statements)
575    }
576
577    /// Builds an expression from a green node.
578    fn build_expression(&self, node: &GreenNode<PythonLanguage>, offset: usize, source: &SourceText) -> Result<Expression, OakError> {
579        match node.kind {
580            PythonElementType::Constant => {
581                let mut current_offset = offset;
582                for child in node.children() {
583                    if let GreenTree::Leaf(leaf) = child {
584                        if !leaf.kind.is_trivia() {
585                            let text = source.get_text_in((current_offset..current_offset + leaf.length as usize).into());
586                            if let Ok(val) = text.parse::<i64>() {
587                                return Ok(Expression::Literal(Literal::Integer(val)));
588                            }
589                            else if let Ok(val) = text.parse::<f64>() {
590                                return Ok(Expression::Literal(Literal::Float(val)));
591                            }
592                            else if text == "True" {
593                                return Ok(Expression::Literal(Literal::Boolean(true)));
594                            }
595                            else if text == "False" {
596                                return Ok(Expression::Literal(Literal::Boolean(false)));
597                            }
598                            else if text == "None" {
599                                return Ok(Expression::Literal(Literal::None));
600                            }
601                            else {
602                                let s = text.to_string();
603                                if s.starts_with('b') || s.starts_with('B') {
604                                    let content = if (s.starts_with("b\"") || s.starts_with("B\"") || s.starts_with("b'")) && (s.ends_with('"') || s.ends_with('\'')) { &s[2..s.len() - 1] } else { &s[1..] };
605                                    return Ok(Expression::Literal(Literal::Bytes(content.as_bytes().to_vec())));
606                                }
607                                let mut s = s;
608                                if (s.starts_with('"') && s.ends_with('"')) || (s.starts_with('\'') && s.ends_with('\'')) {
609                                    s = s[1..s.len() - 1].to_string()
610                                }
611                                return Ok(Expression::Literal(Literal::String(s)));
612                            }
613                        }
614                    }
615                    current_offset += child.len() as usize
616                }
617                Ok(Expression::Name("invalid_constant".to_string()))
618            }
619            PythonElementType::Name => {
620                let mut current_offset = offset;
621                for child in node.children() {
622                    if let GreenTree::Leaf(leaf) = child {
623                        if !leaf.kind.is_trivia() {
624                            let text = source.get_text_in((current_offset..current_offset + leaf.length as usize).into());
625                            return Ok(Expression::Name(text.to_string()));
626                        }
627                    }
628                    current_offset += child.len() as usize;
629                }
630                Ok(Expression::Name("invalid_name".to_string()))
631            }
632            PythonElementType::BinOp => {
633                let mut left = None;
634                let mut operator = None;
635                let mut right = None;
636                let mut current_offset = offset;
637
638                for child in node.children() {
639                    let child_len = child.len() as usize;
640                    match child {
641                        GreenTree::Node(n) => {
642                            if !n.kind.is_trivia() {
643                                let expr = self.build_expression(n, current_offset, source)?;
644                                if left.is_none() {
645                                    left = Some(Box::new(expr));
646                                }
647                                else if right.is_none() {
648                                    right = Some(Box::new(expr));
649                                }
650                            }
651                        }
652                        GreenTree::Leaf(leaf) => {
653                            if !leaf.kind.is_trivia() {
654                                let op = match leaf.kind {
655                                    PythonTokenType::Plus => Some(BinaryOperator::Add),
656                                    PythonTokenType::Minus => Some(BinaryOperator::Sub),
657                                    PythonTokenType::Star => Some(BinaryOperator::Mult),
658                                    PythonTokenType::Slash => Some(BinaryOperator::Div),
659                                    PythonTokenType::Percent => Some(BinaryOperator::Mod),
660                                    PythonTokenType::DoubleStar => Some(BinaryOperator::Pow),
661                                    PythonTokenType::LeftShift => Some(BinaryOperator::LShift),
662                                    PythonTokenType::RightShift => Some(BinaryOperator::RShift),
663                                    PythonTokenType::Pipe => Some(BinaryOperator::BitOr),
664                                    PythonTokenType::Caret => Some(BinaryOperator::BitXor),
665                                    PythonTokenType::Ampersand => Some(BinaryOperator::BitAnd),
666                                    PythonTokenType::DoubleSlash => Some(BinaryOperator::FloorDiv),
667                                    _ => None,
668                                };
669                                if let Some(op) = op {
670                                    operator = Some(op);
671                                }
672                            }
673                        }
674                    }
675                    current_offset += child_len;
676                }
677
678                let l_is = left.is_some();
679                let op_is = operator.is_some();
680                let r_is = right.is_some();
681                if let (Some(l), Some(op), Some(r)) = (left, operator, right) {
682                    Ok(Expression::BinaryOp { left: l, operator: op, right: r })
683                }
684                else {
685                    println!("Warning: Invalid BinOp at {}, left={}, op={}, right={}", offset, l_is, op_is, r_is);
686                    Ok(Expression::Name(format!("invalid_binop_at_{}", offset)))
687                }
688            }
689            PythonElementType::Call => {
690                let mut func = None;
691                let mut args = Vec::new();
692                let mut keywords = Vec::new();
693                let mut current_offset = offset;
694
695                for child in node.children() {
696                    let child_len = child.len() as usize;
697                    if let GreenTree::Node(n) = child {
698                        if !n.kind.is_trivia() {
699                            if func.is_none() {
700                                func = Some(Box::new(self.build_expression(n, current_offset, source)?));
701                            }
702                            else if n.kind == PythonElementType::Keyword {
703                                keywords.push(self.build_keyword(n, current_offset, source)?);
704                            }
705                            else if n.kind == PythonElementType::Starred {
706                                let expr = self.build_starred(n, current_offset, source)?;
707                                args.push(expr);
708                            }
709                            else {
710                                args.push(self.build_expression(n, current_offset, source)?);
711                            }
712                        }
713                    }
714                    current_offset += child_len;
715                }
716
717                if let Some(f) = func { Ok(Expression::Call { func: f, args, keywords }) } else { Ok(Expression::Name("invalid_call".to_string())) }
718            }
719            PythonElementType::Attribute => {
720                let mut value = None;
721                let mut attr = String::new();
722                let mut current_offset = offset;
723
724                for child in node.children() {
725                    let child_len = child.len() as usize;
726                    match child {
727                        GreenTree::Node(n) => {
728                            if !n.kind.is_trivia() {
729                                value = Some(Box::new(self.build_expression(n, current_offset, source)?));
730                            }
731                        }
732                        GreenTree::Leaf(leaf) if leaf.kind == PythonTokenType::Identifier => {
733                            attr = source.get_text_in((current_offset..current_offset + leaf.length as usize).into()).trim().to_string();
734                        }
735                        _ => {}
736                    }
737                    current_offset += child_len;
738                }
739
740                if let Some(v) = value { Ok(Expression::Attribute { value: v, attr }) } else { Ok(Expression::Name("invalid_attribute".to_string())) }
741            }
742            PythonElementType::Expr => {
743                let mut current_offset = offset;
744                for child in node.children() {
745                    let child_len = child.len() as usize;
746                    if let GreenTree::Node(n) = child {
747                        if !n.kind.is_trivia() {
748                            return self.build_expression(n, current_offset, source);
749                        }
750                    }
751                    current_offset += child_len;
752                }
753                Ok(Expression::Name("invalid_expr".to_string()))
754            }
755            PythonElementType::Tuple => {
756                let mut exprs = Vec::new();
757                let mut current_offset = offset;
758                for child in node.children() {
759                    let child_len = child.len() as usize;
760                    if let GreenTree::Node(n) = child {
761                        if !n.kind.is_trivia() {
762                            exprs.push(self.build_expression(n, current_offset, source)?);
763                        }
764                    }
765                    current_offset += child_len;
766                }
767                if exprs.len() == 1 { Ok(exprs.remove(0)) } else { Ok(Expression::Tuple { elts: exprs }) }
768            }
769            PythonElementType::List => {
770                let mut exprs = Vec::new();
771                let mut current_offset = offset;
772                for child in node.children() {
773                    let child_len = child.len() as usize;
774                    if let GreenTree::Node(n) = child {
775                        if !n.kind.is_trivia() {
776                            exprs.push(self.build_expression(n, current_offset, source)?);
777                        }
778                    }
779                    current_offset += child_len;
780                }
781                Ok(Expression::List { elts: exprs })
782            }
783            PythonElementType::Subscript => {
784                let mut value = None;
785                let mut slice = None;
786                let mut current_offset = offset;
787
788                for child in node.children() {
789                    let child_len = child.len() as usize;
790                    if let GreenTree::Node(n) = child {
791                        if !n.kind.is_trivia() {
792                            if value.is_none() {
793                                value = Some(Box::new(self.build_expression(n, current_offset, source)?));
794                            }
795                            else {
796                                slice = Some(Box::new(self.build_expression(n, current_offset, source)?));
797                            }
798                        }
799                    }
800                    current_offset += child_len;
801                }
802
803                if let (Some(v), Some(s)) = (value, slice) { Ok(Expression::Subscript { value: v, slice: s }) } else { Ok(Expression::Name("invalid_subscript".to_string())) }
804            }
805            PythonElementType::JoinedStr => {
806                let mut current_offset = offset;
807                let mut values = Vec::new();
808                for child in node.children() {
809                    let child_len = child.len() as usize;
810                    match child {
811                        GreenTree::Node(n) => {
812                            if !n.kind.is_trivia() {
813                                values.push(self.build_expression(n, current_offset, source)?);
814                            }
815                        }
816                        GreenTree::Leaf(leaf) if leaf.kind == PythonTokenType::FString => {
817                            let text = source.get_text_in((current_offset..current_offset + leaf.length as usize).into());
818                            values.extend(self.parse_fstring_text(&text)?);
819                        }
820                        _ => {}
821                    }
822                    current_offset += child_len;
823                }
824                Ok(Expression::JoinedStr { values })
825            }
826            PythonElementType::FormattedValue => {
827                let mut current_offset = offset;
828                for child in node.children() {
829                    let child_len = child.len() as usize;
830                    if let GreenTree::Node(n) = child {
831                        if !n.kind.is_trivia() {
832                            let value = Box::new(self.build_expression(n, current_offset, source)?);
833                            return Ok(Expression::FormattedValue { value, conversion: 0, format_spec: None });
834                        }
835                    }
836                    current_offset += child_len;
837                }
838                Ok(Expression::Name("invalid_formatted_value".to_string()))
839            }
840            PythonElementType::Yield => {
841                let mut value = None;
842                let mut current_offset = offset;
843                for child in node.children() {
844                    let child_len = child.len() as usize;
845                    if let GreenTree::Node(n) = child {
846                        if !n.kind.is_trivia() {
847                            value = Some(Box::new(self.build_expression(n, current_offset, source)?));
848                        }
849                    }
850                    current_offset += child_len;
851                }
852                Ok(Expression::Yield(value))
853            }
854            PythonElementType::YieldFrom => {
855                let mut value = None;
856                let mut current_offset = offset;
857                for child in node.children() {
858                    let child_len = child.len() as usize;
859                    if let GreenTree::Node(n) = child {
860                        if !n.kind.is_trivia() {
861                            value = Some(Box::new(self.build_expression(n, current_offset, source)?));
862                        }
863                    }
864                    current_offset += child_len;
865                }
866                Ok(Expression::YieldFrom(value.unwrap_or(Box::new(Expression::Name("invalid_yield_from".to_string())))))
867            }
868            PythonElementType::Starred => self.build_starred(node, offset, source),
869            PythonElementType::Await => {
870                let mut value = None;
871                let mut current_offset = offset;
872                for child in node.children() {
873                    let child_len = child.len() as usize;
874                    if let GreenTree::Node(n) = child {
875                        if !n.kind.is_trivia() {
876                            value = Some(Box::new(self.build_expression(n, current_offset, source)?));
877                        }
878                    }
879                    current_offset += child_len;
880                }
881                Ok(Expression::Await(value.unwrap_or(Box::new(Expression::Name("invalid_await".to_string())))))
882            }
883            PythonElementType::Dict => {
884                let mut keys = Vec::new();
885                let mut values = Vec::new();
886                let mut current_offset = offset;
887                for child in node.children() {
888                    let child_len = child.len() as usize;
889                    if let GreenTree::Node(n) = child {
890                        if !n.kind.is_trivia() {
891                            if n.kind == PythonElementType::Starred {
892                                let expr = self.build_starred(n, current_offset, source)?;
893                                keys.push(None);
894                                values.push(expr);
895                            }
896                            else if keys.len() == values.len() {
897                                // We are expecting a key
898                                let expr = self.build_expression(n, current_offset, source)?;
899                                keys.push(Some(expr));
900                            }
901                            else {
902                                // We are expecting a value
903                                let expr = self.build_expression(n, current_offset, source)?;
904                                values.push(expr);
905                            }
906                        }
907                    }
908                    current_offset += child_len;
909                }
910                Ok(Expression::Dict { keys, values })
911            }
912            PythonElementType::Set => {
913                let mut elts = Vec::new();
914                let mut current_offset = offset;
915                for child in node.children() {
916                    let child_len = child.len() as usize;
917                    if let GreenTree::Node(n) = child {
918                        if !n.kind.is_trivia() {
919                            elts.push(self.build_expression(n, current_offset, source)?);
920                        }
921                    }
922                    current_offset += child_len;
923                }
924                Ok(Expression::Set { elts })
925            }
926            PythonElementType::ListComp => {
927                let mut elt = None;
928                let mut generators = Vec::new();
929                let mut current_offset = offset;
930                for child in node.children() {
931                    let child_len = child.len() as usize;
932                    if let GreenTree::Node(n) = child {
933                        if !n.kind.is_trivia() {
934                            if elt.is_none() {
935                                elt = Some(Box::new(self.build_expression(n, current_offset, source)?));
936                            }
937                            else if n.kind == PythonElementType::Comprehension {
938                                generators.push(self.build_comprehension(n, current_offset, source)?);
939                            }
940                        }
941                    }
942                    current_offset += child_len;
943                }
944                Ok(Expression::ListComp { elt: elt.unwrap_or(Box::new(Expression::Name("invalid_elt".to_string()))), generators })
945            }
946            PythonElementType::SetComp => {
947                let mut elt = None;
948                let mut generators = Vec::new();
949                let mut current_offset = offset;
950                for child in node.children() {
951                    let child_len = child.len() as usize;
952                    if let GreenTree::Node(n) = child {
953                        if !n.kind.is_trivia() {
954                            if elt.is_none() {
955                                elt = Some(Box::new(self.build_expression(n, current_offset, source)?));
956                            }
957                            else if n.kind == PythonElementType::Comprehension {
958                                generators.push(self.build_comprehension(n, current_offset, source)?);
959                            }
960                        }
961                    }
962                    current_offset += child_len;
963                }
964                Ok(Expression::SetComp { elt: elt.unwrap_or(Box::new(Expression::Name("invalid_elt".to_string()))), generators })
965            }
966            PythonElementType::DictComp => {
967                let mut key = None;
968                let mut value = None;
969                let mut generators = Vec::new();
970                let mut current_offset = offset;
971                for child in node.children() {
972                    let child_len = child.len() as usize;
973                    if let GreenTree::Node(n) = child {
974                        if !n.kind.is_trivia() {
975                            if key.is_none() {
976                                key = Some(Box::new(self.build_expression(n, current_offset, source)?));
977                            }
978                            else if value.is_none() {
979                                value = Some(Box::new(self.build_expression(n, current_offset, source)?));
980                            }
981                            else if n.kind == PythonElementType::Comprehension {
982                                generators.push(self.build_comprehension(n, current_offset, source)?);
983                            }
984                        }
985                    }
986                    current_offset += child_len;
987                }
988                Ok(Expression::DictComp { key: key.unwrap_or(Box::new(Expression::Name("invalid_key".to_string()))), value: value.unwrap_or(Box::new(Expression::Name("invalid_value".to_string()))), generators })
989            }
990            PythonElementType::GeneratorExp => {
991                let mut elt = None;
992                let mut generators = Vec::new();
993                let mut current_offset = offset;
994                for child in node.children() {
995                    let child_len = child.len() as usize;
996                    if let GreenTree::Node(n) = child {
997                        if !n.kind.is_trivia() {
998                            if elt.is_none() {
999                                elt = Some(Box::new(self.build_expression(n, current_offset, source)?));
1000                            }
1001                            else if n.kind == PythonElementType::Comprehension {
1002                                generators.push(self.build_comprehension(n, current_offset, source)?);
1003                            }
1004                        }
1005                    }
1006                    current_offset += child_len;
1007                }
1008                Ok(Expression::GeneratorExp { elt: elt.unwrap_or(Box::new(Expression::Name("invalid_elt".to_string()))), generators })
1009            }
1010            PythonElementType::Slice => {
1011                let mut lower = None;
1012                let mut upper = None;
1013                let mut step = None;
1014                let mut current_offset = offset;
1015                let mut colon_count = 0;
1016
1017                for child in node.children() {
1018                    let child_len = child.len() as usize;
1019                    match child {
1020                        GreenTree::Leaf(leaf) if leaf.kind == PythonTokenType::Colon => {
1021                            colon_count += 1;
1022                        }
1023                        GreenTree::Node(n) if !n.kind.is_trivia() => {
1024                            let expr = Some(Box::new(self.build_expression(n, current_offset, source)?));
1025                            if colon_count == 0 {
1026                                lower = expr;
1027                            }
1028                            else if colon_count == 1 {
1029                                upper = expr;
1030                            }
1031                            else {
1032                                step = expr;
1033                            }
1034                        }
1035                        _ => {}
1036                    }
1037                    current_offset += child_len;
1038                }
1039                Ok(Expression::Slice { lower, upper, step })
1040            }
1041            _ => Ok(Expression::Name(format!("unsupported_kind_{:?}", node.kind))),
1042        }
1043    }
1044
1045    /// Builds a keyword argument from a green node.
1046    pub(crate) fn build_keyword(&self, node: &GreenNode<PythonLanguage>, offset: usize, source: &SourceText) -> Result<Keyword, OakError> {
1047        let mut arg = None;
1048        let mut value = None;
1049        let mut current_offset = offset;
1050
1051        for child in node.children() {
1052            let child_len = child.len() as usize;
1053            match child {
1054                GreenTree::Leaf(leaf) if leaf.kind == PythonTokenType::Identifier => {
1055                    arg = Some(source.get_text_in((current_offset..current_offset + leaf.length as usize).into()).trim().to_string());
1056                }
1057                GreenTree::Node(n) if !n.kind.is_trivia() => {
1058                    value = Some(self.build_expression(n, current_offset, source)?);
1059                }
1060                _ => {}
1061            }
1062            current_offset += child_len;
1063        }
1064
1065        if let Some(v) = value { Ok(Keyword { arg, value: v }) } else { Err(OakError::custom_error("Invalid keyword".to_string())) }
1066    }
1067
1068    /// Builds a starred expression from a green node.
1069    pub(crate) fn build_starred(&self, node: &GreenNode<PythonLanguage>, offset: usize, source: &SourceText) -> Result<Expression, OakError> {
1070        let mut value = None;
1071        let mut is_double = false;
1072        let mut current_offset = offset;
1073
1074        for child in node.children() {
1075            let child_len = child.len() as usize;
1076            match child {
1077                GreenTree::Leaf(leaf) if leaf.kind == PythonTokenType::Star => {
1078                    is_double = false;
1079                }
1080                GreenTree::Leaf(leaf) if leaf.kind == PythonTokenType::DoubleStar => {
1081                    is_double = true;
1082                }
1083                GreenTree::Node(n) if !n.kind.is_trivia() => {
1084                    value = Some(Box::new(self.build_expression(n, current_offset, source)?));
1085                }
1086                _ => {}
1087            }
1088            current_offset += child_len;
1089        }
1090
1091        if let Some(v) = value { Ok(Expression::Starred { value: v, is_double }) } else { Err(OakError::custom_error("Invalid starred expression".to_string())) }
1092    }
1093
1094    /// Builds an import name (alias) from a green node.
1095    fn build_import_name(&self, node: &GreenNode<PythonLanguage>, offset: usize, source: &SourceText) -> Result<ImportName, OakError> {
1096        let mut name = String::new();
1097        let mut asname = None;
1098        let mut current_offset = offset;
1099
1100        for child in node.children() {
1101            let child_len = child.len() as usize;
1102            match child {
1103                GreenTree::Leaf(leaf) if leaf.kind == PythonTokenType::Identifier => {
1104                    let text = source.get_text_in((current_offset..current_offset + leaf.length as usize).into()).trim().to_string();
1105                    if name.is_empty() {
1106                        name = text;
1107                    }
1108                    else {
1109                        asname = Some(text);
1110                    }
1111                }
1112                _ => {}
1113            }
1114            current_offset += child_len;
1115        }
1116
1117        Ok(ImportName { name, asname })
1118    }
1119
1120    /// Builds a list of parameters from a green node.
1121    fn build_parameters(&self, node: &GreenNode<PythonLanguage>, offset: usize, source: &SourceText) -> Result<Vec<Parameter>, OakError> {
1122        let mut parameters = Vec::new();
1123        let mut current_offset = offset;
1124
1125        for child in node.children() {
1126            let child_len = child.len() as usize;
1127            if let GreenTree::Node(n) = child {
1128                if n.kind == PythonElementType::Arg {
1129                    parameters.push(self.build_parameter(n, current_offset, source)?);
1130                }
1131            }
1132            current_offset += child_len;
1133        }
1134        Ok(parameters)
1135    }
1136
1137    /// Builds a single parameter from a green node.
1138    fn build_parameter(&self, node: &GreenNode<PythonLanguage>, offset: usize, source: &SourceText) -> Result<Parameter, OakError> {
1139        let mut name = String::new();
1140        let mut default = None;
1141        let mut is_vararg = false;
1142        let mut is_kwarg = false;
1143        let mut current_offset = offset;
1144
1145        for child in node.children() {
1146            let child_len = child.len() as usize;
1147            match child {
1148                GreenTree::Leaf(leaf) => {
1149                    if leaf.kind == PythonTokenType::Identifier {
1150                        name = source.get_text_in((current_offset..current_offset + leaf.length as usize).into()).trim().to_string();
1151                    }
1152                    else if leaf.kind == PythonTokenType::Star {
1153                        is_vararg = true;
1154                    }
1155                    else if leaf.kind == PythonTokenType::DoubleStar {
1156                        is_kwarg = true;
1157                    }
1158                }
1159                GreenTree::Node(n) => {
1160                    if !n.kind.is_trivia() {
1161                        default = Some(self.build_expression(n, current_offset, source)?);
1162                    }
1163                }
1164            }
1165            current_offset += child_len;
1166        }
1167
1168        Ok(Parameter { name, annotation: None, default, is_vararg, is_kwarg })
1169    }
1170
1171    /// Builds an exception handler from a green node.
1172    fn build_except_handler(&self, node: &GreenNode<PythonLanguage>, offset: usize, source: &SourceText) -> Result<ExceptHandler, OakError> {
1173        let mut type_ = None;
1174        let mut name = None;
1175        let mut body = Vec::new();
1176        let mut current_offset = offset;
1177
1178        for child in node.children() {
1179            let child_len = child.len() as usize;
1180            match child {
1181                GreenTree::Node(n) => {
1182                    if !n.kind.is_trivia() {
1183                        if n.kind == PythonElementType::Suite {
1184                            body = self.build_suite(n, current_offset, source)?;
1185                        }
1186                        else if type_.is_none() {
1187                            type_ = Some(self.build_expression(n, current_offset, source)?);
1188                        }
1189                    }
1190                }
1191                GreenTree::Leaf(leaf) if leaf.kind == PythonTokenType::Identifier => {
1192                    name = Some(source.get_text_in((current_offset..current_offset + leaf.length as usize).into()).trim().to_string());
1193                }
1194                _ => {}
1195            }
1196            current_offset += child_len;
1197        }
1198        Ok(ExceptHandler { type_, name, body })
1199    }
1200
1201    /// Builds a `with` item from a green node.
1202    fn build_with_item(&self, node: &GreenNode<PythonLanguage>, offset: usize, source: &SourceText) -> Result<WithItem, OakError> {
1203        let mut context_expr = None;
1204        let mut optional_vars = None;
1205        let mut current_offset = offset;
1206
1207        for child in node.children() {
1208            let child_len = child.len() as usize;
1209            if let GreenTree::Node(n) = child {
1210                if !n.kind.is_trivia() {
1211                    let expr = self.build_expression(n, current_offset, source)?;
1212                    if context_expr.is_none() { context_expr = Some(expr) } else { optional_vars = Some(expr) }
1213                }
1214            }
1215            current_offset += child_len
1216        }
1217        Ok(WithItem { context_expr: context_expr.unwrap_or(Expression::Literal(Literal::None)), optional_vars })
1218    }
1219
1220    /// Builds a comprehension from a green node.
1221    fn build_comprehension(&self, node: &GreenNode<PythonLanguage>, offset: usize, source: &SourceText) -> Result<crate::ast::Comprehension, OakError> {
1222        let mut target = None;
1223        let mut iter = None;
1224        let mut ifs = Vec::new();
1225        let mut is_async = false;
1226        let mut current_offset = offset;
1227
1228        for child in node.children() {
1229            let child_len = child.len() as usize;
1230            match child {
1231                GreenTree::Leaf(leaf) if leaf.kind == PythonTokenType::AsyncKeyword => is_async = true,
1232                GreenTree::Node(n) if !n.kind.is_trivia() && n.kind != PythonElementType::ForKeyword && n.kind != PythonElementType::InKeyword && n.kind != PythonElementType::IfKeyword => {
1233                    let expr = self.build_expression(n, current_offset, source)?;
1234                    if target.is_none() {
1235                        target = Some(expr)
1236                    }
1237                    else if iter.is_none() {
1238                        iter = Some(expr)
1239                    }
1240                    else {
1241                        ifs.push(expr)
1242                    }
1243                }
1244                _ => {}
1245            }
1246            current_offset += child_len
1247        }
1248
1249        Ok(crate::ast::Comprehension { target: target.unwrap_or(Expression::Name("invalid_target".to_string())), iter: iter.unwrap_or(Expression::Name("invalid_iter".to_string())), ifs, is_async })
1250    }
1251}
oak_python/builder/mod.rs

oak_python/builder/
mod.rs
