Skip to main content

oak_python/parser/
mod.rs

1pub mod element_type;
2
3use crate::{
4    language::PythonLanguage,
5    lexer::{PythonLexer, token_type::PythonTokenType},
6};
7use oak_core::{
8    OakError,
9    parser::{
10        ParseCache, ParseOutput, Parser, ParserState, parse_with_lexer,
11        pratt::{Associativity, Pratt, PrattParser},
12    },
13    source::{Source, TextEdit},
14    tree::GreenNode,
15};
16
17pub(crate) type State<'a, S> = ParserState<'a, PythonLanguage, S>;
18
19/// Python parser implementation.
20pub struct PythonParser<'config> {
21    /// Reference to the language configuration.
22    pub(crate) config: &'config PythonLanguage,
23}
24
25impl<'config> PythonParser<'config> {
26    /// Creates a new Python parser.
27    pub fn new(config: &'config PythonLanguage) -> Self {
28        Self { config }
29    }
30
31    /// Advances the parser until it reaches a token of the specified kind.
32    fn advance_until<'a, S: Source + ?Sized>(&self, state: &mut State<'a, S>, kind: PythonTokenType) {
33        while state.not_at_end() && !state.at(kind) {
34            state.advance()
35        }
36    }
37
38    /// Skips trivia tokens (whitespace, comments, etc.).
39    fn skip_trivia<'a, S: Source + ?Sized>(&self, state: &mut State<'a, S>) {
40        while state.not_at_end() {
41            if let Some(kind) = state.peek_kind() {
42                if kind.is_trivia() {
43                    state.bump();
44                    continue;
45                }
46            }
47            break;
48        }
49    }
50
51    /// Parses an expression with the given minimum precedence.
52    fn parse_expression<'a, S: Source + ?Sized>(&self, state: &mut State<'a, S>, min_precedence: u8) -> &'a GreenNode<'a, PythonLanguage> {
53        PrattParser::parse(state, min_precedence, self)
54    }
55
56    /// Parses a single statement.
57    pub(crate) fn parse_statement<'a, S: Source + ?Sized>(&self, state: &mut State<'a, S>) -> Result<(), OakError> {
58        use crate::{lexer::PythonTokenType::*, parser::element_type::PythonElementType as ET};
59        let _cp = state.checkpoint();
60        self.skip_trivia(state);
61
62        // Skip leading newlines at top level
63        while state.eat(Newline) {
64            self.skip_trivia(state)
65        }
66
67        if !state.not_at_end() || state.at(Dedent) || state.at(Eof) {
68            return Ok(());
69        }
70
71        if state.at(At) {
72            let cp = state.checkpoint();
73            state.bump();
74            self.parse_expression(state, 0);
75            state.finish_at(cp, crate::parser::element_type::PythonElementType::Decorator);
76            self.skip_trivia(state);
77            state.eat(Newline);
78            return self.parse_statement(state);
79        }
80
81        if state.at(DefKeyword) {
82            state.incremental_node(ET::FunctionDef.into(), |state| self.parse_function_def_body(state))
83        }
84        else if state.at(AsyncKeyword) {
85            state.bump(); // Consume async
86            self.skip_trivia(state);
87            if state.at(DefKeyword) {
88                state.incremental_node(ET::AsyncFunctionDef.into(), |state| self.parse_function_def_body(state))
89            }
90            else if state.at(ForKeyword) {
91                state.incremental_node(ET::AsyncFor.into(), |state| self.parse_for_stmt_body(state))
92            }
93            else if state.at(WithKeyword) {
94                state.incremental_node(ET::AsyncWith.into(), |state| self.parse_with_stmt_body(state))
95            }
96            else {
97                state.incremental_node(ET::Error.into(), |state| {
98                    state.bump();
99                    Ok(())
100                })
101            }
102        }
103        else if state.at(ClassKeyword) {
104            state.incremental_node(ET::ClassDef.into(), |state| self.parse_class_def_body(state))
105        }
106        else if state.at(IfKeyword) {
107            state.incremental_node(ET::If.into(), |state| self.parse_if_stmt_body(state))
108        }
109        else if state.at(WhileKeyword) {
110            state.incremental_node(ET::While.into(), |state| self.parse_while_stmt_body(state))
111        }
112        else if state.at(ForKeyword) {
113            state.incremental_node(ET::For.into(), |state| self.parse_for_stmt_body(state))
114        }
115        else if state.at(TryKeyword) {
116            state.incremental_node(ET::Try.into(), |state| self.parse_try_stmt_body(state))
117        }
118        else if state.at(WithKeyword) {
119            state.incremental_node(ET::With.into(), |state| self.parse_with_stmt_body(state))
120        }
121        else if state.at(RaiseKeyword) {
122            state.incremental_node(ET::Raise.into(), |state| self.parse_raise_stmt_body(state))
123        }
124        else if state.at(AssertKeyword) {
125            state.incremental_node(ET::Assert.into(), |state| self.parse_assert_stmt_body(state))
126        }
127        else if state.at(DelKeyword) {
128            state.incremental_node(ET::Delete.into(), |state| self.parse_del_stmt_body(state))
129        }
130        else if state.at(GlobalKeyword) {
131            state.incremental_node(ET::Global.into(), |state| self.parse_global_stmt_body(state))
132        }
133        else if state.at(NonlocalKeyword) {
134            state.incremental_node(ET::Nonlocal.into(), |state| self.parse_nonlocal_stmt_body(state))
135        }
136        else if state.eat(ReturnKeyword) {
137            let cp = state.checkpoint();
138            self.parse_return_stmt_body(state)?;
139            state.finish_at(cp, ET::Return);
140            state.eat(Newline);
141            Ok(())
142        }
143        else if state.at(ImportKeyword) || state.at(FromKeyword) {
144            state.incremental_node(ET::Import.into(), |state| self.parse_import_stmt_body(state))
145        }
146        else if state.eat(PassKeyword) {
147            state.incremental_node(ET::Pass.into(), |state| {
148                self.skip_trivia(state);
149                state.eat(Newline);
150                Ok(())
151            })
152        }
153        else if state.eat(BreakKeyword) {
154            state.incremental_node(ET::Break.into(), |state| {
155                self.skip_trivia(state);
156                state.eat(Newline);
157                Ok(())
158            })
159        }
160        else if state.eat(ContinueKeyword) {
161            state.incremental_node(ET::Continue.into(), |state| {
162                self.skip_trivia(state);
163                state.eat(Newline);
164                Ok(())
165            })
166        }
167        else {
168            let cp = state.checkpoint();
169            self.parse_expression(state, 0);
170            state.finish_at(cp, ET::Expr);
171            self.skip_trivia(state);
172            state.eat(Newline);
173            Ok(())
174        }
175    }
176
177    /// Parses the body of a function definition.
178    fn parse_function_def_body<'a, S: Source + ?Sized>(&self, state: &mut State<'a, S>) -> Result<(), OakError> {
179        use crate::{lexer::PythonTokenType::*, parser::element_type::PythonElementType as ET};
180        state.expect(DefKeyword).ok();
181        self.skip_trivia(state);
182        if !state.expect(Identifier).is_ok() {
183            // If identifier is missing, we might want to advance to avoid infinite loop
184            // but for now let's just let it be.
185        }
186        self.skip_trivia(state);
187        state.expect(LeftParen).ok();
188        state.incremental_node(ET::Arguments.into(), |state| {
189            while state.not_at_end() && !state.at(RightParen) {
190                self.skip_trivia(state);
191                if state.at(RightParen) {
192                    break;
193                }
194                state.incremental_node(ET::Arg.into(), |state| {
195                    if state.at(Star) {
196                        state.bump();
197                        self.skip_trivia(state);
198                    }
199                    else if state.at(DoubleStar) {
200                        state.bump();
201                        self.skip_trivia(state);
202                    }
203                    state.expect(Identifier).ok();
204                    self.skip_trivia(state);
205                    if state.eat(Colon) {
206                        self.skip_trivia(state);
207                        // Consume until comma or right paren for simple type annotation
208                        while state.not_at_end() && !state.at(Comma) && !state.at(RightParen) && !state.at(Assign) {
209                            state.advance();
210                        }
211                    }
212                    self.skip_trivia(state);
213                    if state.eat(Assign) {
214                        self.skip_trivia(state);
215                        self.parse_expression(state, 0);
216                    }
217                    Ok(())
218                })?;
219                self.skip_trivia(state);
220                if !state.eat(Comma) {
221                    break;
222                }
223            }
224            Ok(())
225        })?;
226        self.skip_trivia(state);
227        state.expect(RightParen).ok();
228        self.skip_trivia(state);
229        if state.eat(Arrow) {
230            self.skip_trivia(state);
231            self.advance_until(state, Colon);
232        }
233        state.expect(Colon).ok();
234        self.parse_suite(state)?;
235        Ok(())
236    }
237
238    /// Parses the body of a class definition.
239    fn parse_class_def_body<'a, S: Source + ?Sized>(&self, state: &mut State<'a, S>) -> Result<(), OakError> {
240        use crate::lexer::PythonTokenType::*;
241        state.expect(ClassKeyword).ok();
242        self.skip_trivia(state);
243        state.expect(Identifier).ok();
244        self.skip_trivia(state);
245        if state.eat(LeftParen) {
246            self.skip_trivia(state);
247            self.advance_until(state, RightParen);
248            state.expect(RightParen).ok();
249        }
250        self.skip_trivia(state);
251        state.expect(Colon).ok();
252        self.parse_suite(state)?;
253        Ok(())
254    }
255
256    /// Parses the body of an if statement.
257    fn parse_if_stmt_body<'a, S: Source + ?Sized>(&self, state: &mut State<'a, S>) -> Result<(), OakError> {
258        use crate::{lexer::PythonTokenType::*, parser::element_type::PythonElementType as ET};
259        state.expect(IfKeyword).ok();
260        self.skip_trivia(state);
261        PrattParser::parse(state, 0, self);
262        self.skip_trivia(state);
263        state.expect(Colon).ok();
264        self.parse_suite(state)?;
265        self.skip_trivia(state);
266
267        // Handle newlines before elif/else
268        let mut lookahead = 0;
269        while let Some(kind) = state.peek_kind_at(lookahead) {
270            if kind == Newline || kind.is_trivia() {
271                lookahead += 1;
272                continue;
273            }
274            if kind == ElifKeyword || kind == ElseKeyword {
275                // Consume the newlines and trivia
276                for _ in 0..lookahead {
277                    state.bump()
278                }
279                break;
280            }
281            break;
282        }
283
284        while state.at(ElifKeyword) {
285            state.incremental_node(ET::If.into(), |state| {
286                state.expect(ElifKeyword).ok();
287                self.skip_trivia(state);
288                self.parse_expression(state, 0);
289                self.skip_trivia(state);
290                state.expect(Colon).ok();
291                self.parse_suite(state)?;
292                self.skip_trivia(state);
293
294                // Peek for more elif/else after newlines
295                let mut lookahead = 0;
296                while let Some(kind) = state.peek_kind_at(lookahead) {
297                    if kind == Newline || kind.is_trivia() {
298                        lookahead += 1;
299                        continue;
300                    }
301                    if kind == ElifKeyword || kind == ElseKeyword {
302                        for _ in 0..lookahead {
303                            state.bump()
304                        }
305                        break;
306                    }
307                    break;
308                }
309                Ok(())
310            })?
311        }
312
313        if state.eat(ElseKeyword) {
314            self.skip_trivia(state);
315            state.expect(Colon).ok();
316            self.parse_suite(state)?
317        }
318        Ok(())
319    }
320
321    /// Parses the body of a while statement.
322    fn parse_while_stmt_body<'a, S: Source + ?Sized>(&self, state: &mut State<'a, S>) -> Result<(), OakError> {
323        use crate::lexer::PythonTokenType::*;
324        state.expect(WhileKeyword).ok();
325        self.skip_trivia(state);
326        PrattParser::parse(state, 0, self);
327        self.skip_trivia(state);
328        state.expect(Colon).ok();
329        self.parse_suite(state)?;
330        self.skip_trivia(state);
331
332        // Handle newlines before else
333        let mut lookahead = 0;
334        while let Some(kind) = state.peek_kind_at(lookahead) {
335            if kind == Newline || kind.is_trivia() {
336                lookahead += 1;
337                continue;
338            }
339            if kind == ElseKeyword {
340                for _ in 0..lookahead {
341                    state.bump()
342                }
343                break;
344            }
345            break;
346        }
347
348        if state.eat(ElseKeyword) {
349            self.skip_trivia(state);
350            state.expect(Colon).ok();
351            self.parse_suite(state)?
352        }
353        Ok(())
354    }
355
356    /// Parses the body of a for statement.
357    fn parse_for_stmt_body<'a, S: Source + ?Sized>(&self, state: &mut State<'a, S>) -> Result<(), OakError> {
358        use crate::lexer::PythonTokenType::*;
359        state.expect(ForKeyword).ok();
360        self.skip_trivia(state);
361
362        // Use a higher precedence to stop at 'in'
363        let target_cp = state.checkpoint();
364        self.parse_expression(state, 6);
365
366        if state.at(Comma) {
367            while state.eat(Comma) {
368                self.skip_trivia(state);
369                if state.at(InKeyword) {
370                    break;
371                }
372                self.parse_expression(state, 6);
373                self.skip_trivia(state);
374            }
375            state.finish_at(target_cp, crate::parser::element_type::PythonElementType::Tuple);
376        }
377
378        self.skip_trivia(state);
379        state.expect(InKeyword).ok();
380        self.skip_trivia(state);
381        PrattParser::parse(state, 0, self);
382        self.skip_trivia(state);
383        state.expect(Colon).ok();
384        self.parse_suite(state)?;
385        self.skip_trivia(state);
386        if state.eat(ElseKeyword) {
387            self.skip_trivia(state);
388            state.expect(Colon).ok();
389            self.parse_suite(state)?
390        }
391        Ok(())
392    }
393
394    /// Parses the body of a try statement.
395    fn parse_try_stmt_body<'a, S: Source + ?Sized>(&self, state: &mut State<'a, S>) -> Result<(), OakError> {
396        use crate::{lexer::PythonTokenType::*, parser::element_type::PythonElementType as ET};
397        state.expect(TryKeyword).ok();
398        self.skip_trivia(state);
399        state.expect(Colon).ok();
400        self.parse_suite(state)?;
401        self.skip_trivia(state);
402
403        // Handle newlines before except/else/finally
404        let mut lookahead = 0;
405        while let Some(kind) = state.peek_kind_at(lookahead) {
406            if kind == Newline || kind.is_trivia() {
407                lookahead += 1;
408                continue;
409            }
410            if kind == ExceptKeyword || kind == ElseKeyword || kind == FinallyKeyword {
411                for _ in 0..lookahead {
412                    state.bump()
413                }
414                break;
415            }
416            break;
417        }
418
419        while state.at(ExceptKeyword) {
420            state.incremental_node(ET::ExceptHandler.into(), |state| {
421                state.expect(ExceptKeyword).ok();
422                self.skip_trivia(state);
423                if !state.at(Colon) {
424                    self.parse_expression(state, 0); // Exception type
425                    self.skip_trivia(state);
426                    if state.eat(AsKeyword) {
427                        self.skip_trivia(state);
428                        state.expect(Identifier).ok();
429                        self.skip_trivia(state)
430                    }
431                }
432                self.skip_trivia(state);
433                state.expect(Colon).ok();
434                self.parse_suite(state)?;
435                self.skip_trivia(state);
436
437                // Peek for more except/else/finally after newlines
438                let mut lookahead = 0;
439                while let Some(kind) = state.peek_kind_at(lookahead) {
440                    if kind == Newline || kind.is_trivia() {
441                        lookahead += 1;
442                        continue;
443                    }
444                    if kind == ExceptKeyword || kind == ElseKeyword || kind == FinallyKeyword {
445                        for _ in 0..lookahead {
446                            state.bump()
447                        }
448                        break;
449                    }
450                    break;
451                }
452                Ok(())
453            })?
454        }
455
456        if state.eat(ElseKeyword) {
457            self.skip_trivia(state);
458            state.expect(Colon).ok();
459            self.parse_suite(state)?;
460            self.skip_trivia(state);
461
462            // Peek for finally after newlines
463            let mut lookahead = 0;
464            while let Some(kind) = state.peek_kind_at(lookahead) {
465                if kind == Newline || kind.is_trivia() {
466                    lookahead += 1;
467                    continue;
468                }
469                if kind == FinallyKeyword {
470                    for _ in 0..lookahead {
471                        state.bump()
472                    }
473                    break;
474                }
475                break;
476            }
477        }
478
479        if state.eat(FinallyKeyword) {
480            self.skip_trivia(state);
481            state.expect(Colon).ok();
482            self.parse_suite(state)?
483        }
484        Ok(())
485    }
486
487    /// Parses the body of a with statement.
488    fn parse_with_stmt_body<'a, S: Source + ?Sized>(&self, state: &mut State<'a, S>) -> Result<(), OakError> {
489        use crate::{lexer::PythonTokenType::*, parser::element_type::PythonElementType as ET};
490        state.expect(WithKeyword).ok();
491        self.skip_trivia(state);
492        loop {
493            state.incremental_node(ET::WithItem.into(), |state| {
494                PrattParser::parse(state, 0, self);
495                self.skip_trivia(state);
496                if state.eat(AsKeyword) {
497                    self.skip_trivia(state);
498                    PrattParser::parse(state, 0, self);
499                    self.skip_trivia(state)
500                }
501                Ok(())
502            })?;
503            if !state.eat(Comma) {
504                break;
505            }
506            self.skip_trivia(state)
507        }
508        state.expect(Colon).ok();
509        self.parse_suite(state)?;
510        Ok(())
511    }
512
513    /// Parses the body of a raise statement.
514    fn parse_raise_stmt_body<'a, S: Source + ?Sized>(&self, state: &mut State<'a, S>) -> Result<(), OakError> {
515        use crate::lexer::PythonTokenType::*;
516        state.expect(RaiseKeyword).ok();
517        self.skip_trivia(state);
518        if !state.at(Newline) && !state.at(Semicolon) && state.not_at_end() {
519            PrattParser::parse(state, 0, self);
520            self.skip_trivia(state);
521            if state.eat(FromKeyword) {
522                self.skip_trivia(state);
523                PrattParser::parse(state, 0, self);
524            }
525        }
526        self.skip_trivia(state);
527        state.eat(Newline);
528        Ok(())
529    }
530
531    /// Parses the body of an assert statement.
532    fn parse_assert_stmt_body<'a, S: Source + ?Sized>(&self, state: &mut State<'a, S>) -> Result<(), OakError> {
533        use crate::lexer::PythonTokenType::*;
534        state.expect(AssertKeyword).ok();
535        self.skip_trivia(state);
536        PrattParser::parse(state, 0, self);
537        self.skip_trivia(state);
538        if state.eat(Comma) {
539            self.skip_trivia(state);
540            PrattParser::parse(state, 0, self);
541        }
542        self.skip_trivia(state);
543        state.eat(Newline);
544        Ok(())
545    }
546
547    /// Parses the body of a delete statement.
548    fn parse_del_stmt_body<'a, S: Source + ?Sized>(&self, state: &mut State<'a, S>) -> Result<(), OakError> {
549        use crate::lexer::PythonTokenType::*;
550        state.expect(DelKeyword).ok();
551        self.skip_trivia(state);
552        PrattParser::parse(state, 0, self);
553        self.skip_trivia(state);
554        state.eat(Newline);
555        Ok(())
556    }
557
558    /// Parses the body of a global statement.
559    fn parse_global_stmt_body<'a, S: Source + ?Sized>(&self, state: &mut State<'a, S>) -> Result<(), OakError> {
560        use crate::lexer::PythonTokenType::*;
561        state.expect(GlobalKeyword).ok();
562        self.skip_trivia(state);
563        loop {
564            state.expect(Identifier).ok();
565            self.skip_trivia(state);
566            if !state.eat(Comma) {
567                break;
568            }
569            self.skip_trivia(state)
570        }
571        state.eat(Newline);
572        Ok(())
573    }
574
575    /// Parses the body of a nonlocal statement.
576    fn parse_nonlocal_stmt_body<'a, S: Source + ?Sized>(&self, state: &mut State<'a, S>) -> Result<(), OakError> {
577        use crate::lexer::PythonTokenType::*;
578        state.expect(NonlocalKeyword).ok();
579        self.skip_trivia(state);
580        loop {
581            state.expect(Identifier).ok();
582            self.skip_trivia(state);
583            if !state.eat(Comma) {
584                break;
585            }
586            self.skip_trivia(state)
587        }
588        state.eat(Newline);
589        Ok(())
590    }
591
592    /// Parses the body of a return statement.
593    fn parse_return_stmt_body<'a, S: Source + ?Sized>(&self, state: &mut State<'a, S>) -> Result<(), OakError> {
594        use crate::lexer::PythonTokenType::*;
595        self.skip_trivia(state);
596        if state.not_at_end() && !state.at(Newline) && !state.at(Semicolon) {
597            self.parse_expression(state, 0);
598        }
599        Ok(())
600    }
601
602    /// Parses the body of an import statement.
603    fn parse_import_stmt_body<'a, S: Source + ?Sized>(&self, state: &mut State<'a, S>) -> Result<(), OakError> {
604        use crate::lexer::PythonTokenType::*;
605        if state.at(ImportKeyword) {
606            state.bump();
607            self.skip_trivia(state);
608            loop {
609                state
610                    .incremental_node(crate::parser::element_type::PythonElementType::Alias.into(), |state| {
611                        state.expect(Identifier).ok();
612                        self.skip_trivia(state);
613                        if state.eat(AsKeyword) {
614                            self.skip_trivia(state);
615                            state.expect(Identifier).ok();
616                            self.skip_trivia(state)
617                        }
618                        Ok(())
619                    })
620                    .ok();
621                if !state.eat(Comma) {
622                    break;
623                }
624                self.skip_trivia(state)
625            }
626        }
627        else if state.at(FromKeyword) {
628            state.bump();
629            self.skip_trivia(state);
630            // Module name (can be dotted)
631            while state.at(Identifier) || state.at(Dot) {
632                state.bump();
633                self.skip_trivia(state)
634            }
635            state.expect(ImportKeyword).ok();
636            self.skip_trivia(state);
637            if state.eat(Star) {
638                // from module import *
639            }
640            else {
641                loop {
642                    state
643                        .incremental_node(crate::parser::element_type::PythonElementType::Alias.into(), |state| {
644                            state.expect(Identifier).ok();
645                            self.skip_trivia(state);
646                            if state.eat(AsKeyword) {
647                                self.skip_trivia(state);
648                                state.expect(Identifier).ok();
649                                self.skip_trivia(state)
650                            }
651                            Ok(())
652                        })
653                        .ok();
654                    if !state.eat(Comma) {
655                        break;
656                    }
657                    self.skip_trivia(state)
658                }
659            }
660        }
661        state.eat(Newline);
662        Ok(())
663    }
664
665    /// Parses a suite (a block of statements).
666    fn parse_suite<'a, S: Source + ?Sized>(&self, state: &mut State<'a, S>) -> Result<(), OakError> {
667        use crate::lexer::PythonTokenType::*;
668        let cp = state.checkpoint();
669        self.skip_trivia(state);
670        if state.eat(Newline) {
671            self.skip_trivia(state);
672            state.expect(Indent).ok();
673            while state.not_at_end() && !state.at(Dedent) {
674                self.parse_statement(state)?;
675                self.skip_trivia(state);
676            }
677            state.expect(Dedent).ok();
678            ()
679        }
680        else {
681            self.parse_statement(state)?
682        }
683        state.finish_at(cp, crate::parser::element_type::PythonElementType::Suite);
684        Ok(())
685    }
686
687    /// Internal method to parse the root module.
688    fn parse_root_internal<'a, S: Source + ?Sized>(&self, state: &mut State<'a, S>) -> Result<&'a GreenNode<'a, PythonLanguage>, OakError> {
689        let checkpoint = state.checkpoint();
690
691        while state.not_at_end() && !state.at(PythonTokenType::Eof) {
692            let start_index = state.tokens.index();
693            self.parse_statement(state)?;
694
695            // Safety: if parse_statement didn't advance, we must advance to avoid infinite loop
696            if state.tokens.index() == start_index && state.not_at_end() {
697                state.bump()
698            }
699        }
700
701        if state.at(PythonTokenType::Eof) {
702            state.bump()
703        }
704
705        Ok(state.finish_at(checkpoint, crate::parser::element_type::PythonElementType::ExpressionModule))
706    }
707
708    /// Parses a comprehension (list, set, dict, or generator).
709    fn parse_comprehension<'a, S: Source + ?Sized>(&self, state: &mut State<'a, S>) {
710        use crate::{lexer::PythonTokenType::*, parser::element_type::PythonElementType as ET};
711        while state.at(ForKeyword) || state.at(AsyncKeyword) {
712            state
713                .incremental_node(ET::Comprehension.into(), |state| {
714                    state.eat(AsyncKeyword);
715                    self.skip_trivia(state);
716                    state.expect(ForKeyword).ok();
717                    self.skip_trivia(state);
718                    self.parse_expression(state, 6); // target
719                    self.skip_trivia(state);
720                    state.expect(InKeyword).ok();
721                    self.skip_trivia(state);
722                    self.parse_expression(state, 6); // iter
723                    self.skip_trivia(state);
724                    while state.at(IfKeyword) {
725                        state.bump();
726                        self.skip_trivia(state);
727                        self.parse_expression(state, 6); // if condition
728                        self.skip_trivia(state)
729                    }
730                    Ok(())
731                })
732                .ok();
733            self.skip_trivia(state)
734        }
735    }
736}
737
738impl<'config> Pratt<PythonLanguage> for PythonParser<'config> {
739    fn primary<'a, S: Source + ?Sized>(&self, state: &mut State<'a, S>) -> &'a GreenNode<'a, PythonLanguage> {
740        use crate::{lexer::PythonTokenType::*, parser::element_type::PythonElementType as ET};
741        self.skip_trivia(state);
742
743        println!("primary at {}: {:?}", state.current_offset(), state.peek_kind());
744
745        let cp = state.checkpoint();
746        let kind = state.peek_kind();
747        match kind {
748            Some(Identifier) => {
749                state.bump();
750                state.finish_at(cp, ET::Name)
751            }
752            Some(Number) | Some(String) | Some(Bytes) | Some(TrueKeyword) | Some(FalseKeyword) | Some(NoneKeyword) => {
753                state.bump();
754                state.finish_at(cp, ET::Constant)
755            }
756            Some(FString) => {
757                state.bump();
758                state.finish_at(cp, ET::JoinedStr)
759            }
760            Some(LeftParen) => {
761                state.bump();
762                self.skip_trivia(state);
763                if state.eat(RightParen) {
764                    state.finish_at(cp, ET::Tuple)
765                }
766                else {
767                    self.parse_expression(state, 0);
768                    self.skip_trivia(state);
769                    if state.at(Comma) {
770                        while state.eat(Comma) {
771                            self.skip_trivia(state);
772                            if state.at(RightParen) {
773                                break;
774                            }
775                            self.parse_expression(state, 0);
776                            self.skip_trivia(state);
777                        }
778                        state.expect(RightParen).ok();
779                        state.finish_at(cp, ET::Tuple)
780                    }
781                    else {
782                        state.expect(RightParen).ok();
783                        state.finish_at(cp, ET::Expr)
784                    }
785                }
786            }
787            Some(LeftBracket) => {
788                state.bump();
789                self.skip_trivia(state);
790                if state.eat(RightBracket) {
791                    state.finish_at(cp, ET::List)
792                }
793                else {
794                    self.parse_expression(state, 0);
795                    self.skip_trivia(state);
796                    if state.at(ForKeyword) || state.at(AsyncKeyword) {
797                        self.parse_comprehension(state);
798                        state.expect(RightBracket).ok();
799                        state.finish_at(cp, ET::ListComp)
800                    }
801                    else {
802                        while state.eat(Comma) {
803                            self.skip_trivia(state);
804                            if state.at(RightBracket) {
805                                break;
806                            }
807                            self.parse_expression(state, 0);
808                            self.skip_trivia(state);
809                        }
810                        state.expect(RightBracket).ok();
811                        state.finish_at(cp, ET::List)
812                    }
813                }
814            }
815            Some(LeftBrace) => {
816                state.bump();
817                self.skip_trivia(state);
818                if state.eat(RightBrace) {
819                    state.finish_at(cp, ET::Dict)
820                }
821                else {
822                    self.parse_expression(state, 0);
823                    self.skip_trivia(state);
824                    if state.eat(Colon) {
825                        self.skip_trivia(state);
826                        self.parse_expression(state, 0);
827                        self.skip_trivia(state);
828                        if state.at(ForKeyword) || state.at(AsyncKeyword) {
829                            self.parse_comprehension(state);
830                            state.expect(RightBrace).ok();
831                            state.finish_at(cp, ET::DictComp)
832                        }
833                        else {
834                            while state.eat(Comma) {
835                                self.skip_trivia(state);
836                                if state.at(RightBrace) {
837                                    break;
838                                }
839                                self.parse_expression(state, 0);
840                                self.skip_trivia(state);
841                                state.expect(Colon).ok();
842                                self.skip_trivia(state);
843                                self.parse_expression(state, 0);
844                                self.skip_trivia(state);
845                            }
846                            state.expect(RightBrace).ok();
847                            state.finish_at(cp, ET::Dict)
848                        }
849                    }
850                    else if state.at(ForKeyword) || state.at(AsyncKeyword) {
851                        self.parse_comprehension(state);
852                        state.expect(RightBrace).ok();
853                        state.finish_at(cp, ET::SetComp)
854                    }
855                    else {
856                        while state.eat(Comma) {
857                            self.skip_trivia(state);
858                            if state.at(RightBrace) {
859                                break;
860                            }
861                            self.parse_expression(state, 0);
862                            self.skip_trivia(state);
863                        }
864                        state.expect(RightBrace).ok();
865                        state.finish_at(cp, ET::Set)
866                    }
867                }
868            }
869            Some(Plus) | Some(Minus) | Some(Tilde) | Some(NotKeyword) => {
870                state.bump();
871                self.skip_trivia(state);
872                self.parse_expression(state, 12);
873                state.finish_at(cp, ET::UnaryOp)
874            }
875            Some(LambdaKeyword) => {
876                state.bump();
877                self.skip_trivia(state);
878                state
879                    .incremental_node(ET::Arguments.into(), |state| {
880                        while state.not_at_end() && !state.at(Colon) {
881                            state.incremental_node(ET::Arg.into(), |state| {
882                                state.expect(Identifier).ok();
883                                self.skip_trivia(state);
884                                if state.eat(Assign) {
885                                    self.skip_trivia(state);
886                                    self.parse_expression(state, 0);
887                                }
888                                Ok(())
889                            })?;
890                            self.skip_trivia(state);
891                            if !state.eat(Comma) {
892                                break;
893                            }
894                            self.skip_trivia(state);
895                        }
896                        Ok(())
897                    })
898                    .ok();
899                state.expect(Colon).ok();
900                self.skip_trivia(state);
901                self.parse_expression(state, 0);
902                state.finish_at(cp, ET::Lambda)
903            }
904            Some(YieldKeyword) => {
905                state.bump();
906                self.skip_trivia(state);
907                if state.eat(FromKeyword) {
908                    self.skip_trivia(state);
909                    self.parse_expression(state, 0);
910                    state.finish_at(cp, ET::YieldFrom)
911                }
912                else {
913                    if !state.at(Newline) && !state.at(RightParen) && !state.at(RightBracket) && !state.at(RightBrace) && !state.at(Comma) && !state.at(Colon) {
914                        self.parse_expression(state, 0);
915                    }
916                    state.finish_at(cp, ET::Yield)
917                }
918            }
919            _ => {
920                state.bump();
921                state.finish_at(cp, ET::Error)
922            }
923        }
924    }
925
926    fn prefix<'a, S: Source + ?Sized>(&self, state: &mut State<'a, S>) -> &'a GreenNode<'a, PythonLanguage> {
927        use crate::lexer::token_type::PythonTokenType;
928        self.skip_trivia(state);
929        let kind = state.peek_kind().expect("Expected token in prefix");
930        match kind {
931            PythonTokenType::Plus | PythonTokenType::Minus | PythonTokenType::Tilde | PythonTokenType::NotKeyword => {
932                let cp = state.checkpoint();
933                state.expect(kind).ok();
934                PrattParser::parse(state, 14, self);
935                state.finish_at(cp, crate::parser::element_type::PythonElementType::UnaryOp)
936            }
937            _ => self.primary(state),
938        }
939    }
940
941    fn infix<'a, S: Source + ?Sized>(&self, state: &mut State<'a, S>, _left: &'a GreenNode<'a, PythonLanguage>, min_precedence: u8) -> Option<&'a GreenNode<'a, PythonLanguage>> {
942        use crate::{lexer::PythonTokenType::*, parser::element_type::PythonElementType as ET};
943
944        // Peek kind without consuming trivia yet
945        let mut lookahead = 0;
946        let mut kind = None;
947        while let Some(k) = state.peek_kind_at(lookahead) {
948            if k.is_trivia() {
949                lookahead += 1;
950                continue;
951            }
952            kind = Some(k);
953            break;
954        }
955
956        let kind = kind?;
957        let (prec, assoc) = match kind {
958            PythonTokenType::Assign
959            | PythonTokenType::PlusAssign
960            | PythonTokenType::MinusAssign
961            | PythonTokenType::StarAssign
962            | PythonTokenType::DoubleStarAssign
963            | PythonTokenType::SlashAssign
964            | PythonTokenType::DoubleSlashAssign
965            | PythonTokenType::PercentAssign
966            | PythonTokenType::AtAssign
967            | PythonTokenType::AmpersandAssign
968            | PythonTokenType::PipeAssign
969            | PythonTokenType::CaretAssign
970            | PythonTokenType::LeftShiftAssign
971            | PythonTokenType::RightShiftAssign => (1, Associativity::Right),
972            PythonTokenType::OrKeyword => (2, Associativity::Left),
973            PythonTokenType::AndKeyword => (3, Associativity::Left),
974            PythonTokenType::NotKeyword => (4, Associativity::Left),
975            PythonTokenType::Less | PythonTokenType::Greater | PythonTokenType::Equal | PythonTokenType::NotEqual | PythonTokenType::LessEqual | PythonTokenType::GreaterEqual | PythonTokenType::InKeyword | PythonTokenType::IsKeyword => {
976                (5, Associativity::Left)
977            }
978            PythonTokenType::Pipe => (6, Associativity::Left),
979            PythonTokenType::Caret => (7, Associativity::Left),
980            PythonTokenType::Ampersand => (8, Associativity::Left),
981            PythonTokenType::LeftShift | PythonTokenType::RightShift => (9, Associativity::Left),
982            PythonTokenType::Plus | PythonTokenType::Minus => (10, Associativity::Left),
983            PythonTokenType::Star | PythonTokenType::Slash | PythonTokenType::DoubleSlash | PythonTokenType::Percent | PythonTokenType::At => (11, Associativity::Left),
984            PythonTokenType::DoubleStar => (13, Associativity::Right),
985            PythonTokenType::Dot | PythonTokenType::LeftParen | PythonTokenType::LeftBracket => (15, Associativity::Left),
986            _ => return None,
987        };
988
989        if prec < min_precedence {
990            return None;
991        }
992
993        let cp = state.checkpoint_before(_left);
994        for _ in 0..lookahead {
995            state.bump();
996        }
997
998        match kind {
999            LeftParen => {
1000                state.bump(); // Consume LeftParen
1001
1002                while state.not_at_end() && !state.at(RightParen) {
1003                    self.skip_trivia(state);
1004                    if state.at(RightParen) {
1005                        break;
1006                    }
1007
1008                    let mut lookahead = 0;
1009                    let mut is_keyword = false;
1010                    if state.at(Identifier) {
1011                        lookahead += 1;
1012                        while let Some(k) = state.peek_kind_at(lookahead) {
1013                            if k.is_trivia() {
1014                                lookahead += 1;
1015                                continue;
1016                            }
1017                            if k == Assign {
1018                                is_keyword = true;
1019                            }
1020                            break;
1021                        }
1022                    }
1023
1024                    if is_keyword {
1025                        state
1026                            .incremental_node(ET::Keyword.into(), |state| {
1027                                state.expect(Identifier).ok();
1028                                self.skip_trivia(state);
1029                                state.expect(Assign).ok();
1030                                self.skip_trivia(state);
1031                                self.parse_expression(state, 0);
1032                                Ok(())
1033                            })
1034                            .ok()?;
1035                    }
1036                    else if state.at(DoubleStar) {
1037                        state
1038                            .incremental_node(ET::Keyword.into(), |state| {
1039                                state.bump(); // DoubleStar
1040                                self.skip_trivia(state);
1041                                self.parse_expression(state, 0);
1042                                Ok(())
1043                            })
1044                            .ok()?;
1045                    }
1046                    else if state.at(Star) {
1047                        state
1048                            .incremental_node(ET::Starred.into(), |state| {
1049                                state.bump(); // Star
1050                                self.skip_trivia(state);
1051                                self.parse_expression(state, 0);
1052                                Ok(())
1053                            })
1054                            .ok()?;
1055                    }
1056                    else {
1057                        self.parse_expression(state, 0);
1058                    }
1059
1060                    self.skip_trivia(state);
1061                    if !state.eat(Comma) {
1062                        break;
1063                    }
1064                }
1065
1066                self.skip_trivia(state);
1067                state.expect(RightParen).ok();
1068                Some(state.finish_at(cp, ET::Call))
1069            }
1070            LeftBracket => {
1071                state.bump(); // Consume LeftBracket
1072
1073                let cp_slice = state.checkpoint();
1074                let mut is_slice = false;
1075
1076                if state.at(Colon) {
1077                    is_slice = true;
1078                }
1079                else {
1080                    self.parse_expression(state, 0);
1081                    self.skip_trivia(state);
1082                    if state.at(Colon) {
1083                        is_slice = true;
1084                    }
1085                }
1086
1087                if is_slice {
1088                    // It's a slice
1089                    while state.at(Colon) {
1090                        state.bump();
1091                        self.skip_trivia(state);
1092                        if !state.at(Colon) && !state.at(RightBracket) {
1093                            self.parse_expression(state, 0);
1094                            self.skip_trivia(state);
1095                        }
1096                    }
1097                    state.finish_at(cp_slice, ET::Slice);
1098                }
1099
1100                state.expect(RightBracket).ok();
1101                Some(state.finish_at(cp, ET::Subscript))
1102            }
1103            Dot => {
1104                state.bump(); // Consume Dot
1105                self.skip_trivia(state);
1106                state.expect(Identifier).ok();
1107                Some(state.finish_at(cp, ET::Attribute))
1108            }
1109            _ => {
1110                let result_kind = if prec == 1 {
1111                    ET::AssignStmt
1112                }
1113                else if prec <= 3 {
1114                    ET::BoolOp
1115                }
1116                else if prec == 5 {
1117                    ET::Compare
1118                }
1119                else {
1120                    ET::BinOp
1121                };
1122
1123                state.bump(); // Consume operator
1124
1125                let next_prec = match assoc {
1126                    Associativity::Left => prec + 1,
1127                    Associativity::Right => prec,
1128                    Associativity::None => prec + 1,
1129                };
1130
1131                self.parse_expression(state, next_prec);
1132                Some(state.finish_at(cp, result_kind.into()))
1133            }
1134        }
1135    }
1136}
1137
1138impl<'config> Parser<PythonLanguage> for PythonParser<'config> {
1139    fn parse<'a, S: Source + ?Sized>(&self, text: &'a S, edits: &[TextEdit], cache: &'a mut impl ParseCache<PythonLanguage>) -> ParseOutput<'a, PythonLanguage> {
1140        let lexer = PythonLexer::new(self.config);
1141        parse_with_lexer(&lexer, text, edits, cache, |state| self.parse_root_internal(state))
1142    }
1143}