Update parser API references, expose Program to linter (#11505)

## Summary

This PR updates the parser API references per the updated function
signatures, especially the return type. Additionally, it exposes the
`Program` to the linter by using the new API and also storing it on the
`Checker` to be used by the rules.

Most of the changes are in test functions. There are only a handful of
usages in the non-test source code.

crates/ruff_benchmark/benches/formatter.rs

@@ -5,9 +5,7 @@ use ruff_benchmark::criterion::{
 };
 use ruff_benchmark::{TestCase, TestFile, TestFileDownloadError};
 use ruff_python_formatter::{format_module_ast, PreviewMode, PyFormatOptions};
-use ruff_python_index::CommentRangesBuilder;
-use ruff_python_parser::lexer::lex;
-use ruff_python_parser::{allocate_tokens_vec, parse_tokens, Mode};
+use ruff_python_parser::{parse, Mode};
 
 #[cfg(target_os = "windows")]
 #[global_allocator]
@@ -52,28 +50,20 @@ fn benchmark_formatter(criterion: &mut Criterion) {
             BenchmarkId::from_parameter(case.name()),
             &case,
             |b, case| {
-                let mut tokens = allocate_tokens_vec(case.code());
-                let mut comment_ranges = CommentRangesBuilder::default();
-
-                for result in lex(case.code(), Mode::Module) {
-                    let (token, range) = result.expect("Input to be a valid python program.");
-
-                    comment_ranges.visit_token(&token, range);
-                    tokens.push(Ok((token, range)));
-                }
-
-                let comment_ranges = comment_ranges.finish();
-
-                // Parse the AST.
-                let module = parse_tokens(tokens, case.code(), Mode::Module)
-                    .expect("Input to be a valid python program");
+                // Parse the source.
+                let program =
+                    parse(case.code(), Mode::Module).expect("Input should be a valid Python code");
 
                 b.iter(|| {
                     let options = PyFormatOptions::from_extension(Path::new(case.name()))
                         .with_preview(PreviewMode::Enabled);
-                    let formatted =
-                        format_module_ast(&module, &comment_ranges, case.code(), options)
-                            .expect("Formatting to succeed");
+                    let formatted = format_module_ast(
+                        program.syntax(),
+                        program.comment_ranges(),
+                        case.code(),
+                        options,
+                    )
+                    .expect("Formatting to succeed");
 
                     formatted.print().expect("Printing to succeed")
                 });

crates/ruff_benchmark/benches/linter.rs

@@ -10,7 +10,7 @@ use ruff_linter::settings::{flags, LinterSettings};
 use ruff_linter::source_kind::SourceKind;
 use ruff_linter::{registry::Rule, RuleSelector};
 use ruff_python_ast::PySourceType;
-use ruff_python_parser::{parse_program_tokens, tokenize, Mode};
+use ruff_python_parser::{parse_module, Mode};
 
 #[cfg(target_os = "windows")]
 #[global_allocator]
@@ -54,11 +54,9 @@ fn benchmark_linter(mut group: BenchmarkGroup, settings: &LinterSettings) {
             BenchmarkId::from_parameter(case.name()),
             &case,
             |b, case| {
-                // Tokenize the source.
-                let tokens = tokenize(case.code(), Mode::Module);
-
                 // Parse the source.
-                let ast = parse_program_tokens(tokens.clone(), case.code(), false).unwrap();
+                let program =
+                    parse_module(case.code()).expect("Input should be a valid Python code");
 
                 b.iter(|| {
                     let path = case.path();
@@ -69,10 +67,7 @@ fn benchmark_linter(mut group: BenchmarkGroup, settings: &LinterSettings) {
                         flags::Noqa::Enabled,
                         &SourceKind::Python(case.code().to_string()),
                         PySourceType::from(path.as_path()),
-                        ParseSource::Precomputed {
-                            tokens: &tokens,
-                            ast: &ast,
-                        },
+                        ParseSource::Precomputed(program),
                     );
 
                     // Assert that file contains no parse errors

crates/ruff_benchmark/benches/parser.rs

@@ -4,7 +4,7 @@ use ruff_benchmark::criterion::{
 use ruff_benchmark::{TestCase, TestFile, TestFileDownloadError};
 use ruff_python_ast::statement_visitor::{walk_stmt, StatementVisitor};
 use ruff_python_ast::Stmt;
-use ruff_python_parser::parse_suite;
+use ruff_python_parser::{parse_module, Mode};
 
 #[cfg(target_os = "windows")]
 #[global_allocator]
@@ -60,7 +60,9 @@ fn benchmark_parser(criterion: &mut Criterion<WallTime>) {
             &case,
             |b, case| {
                 b.iter(|| {
-                    let parsed = parse_suite(case.code()).unwrap();
+                    let parsed = parse_module(case.code())
+                        .expect("Input should be a valid Python code")
+                        .into_suite();
 
                     let mut visitor = CountVisitor { count: 0 };
                     visitor.visit_body(&parsed);

crates/ruff_dev/src/print_ast.rs

@@ -24,7 +24,7 @@ pub(crate) fn main(args: &Args) -> Result<()> {
             args.file.display()
         )
     })?;
-    let python_ast = parse(source_kind.source_code(), source_type.as_mode())?;
+    let python_ast = parse(source_kind.source_code(), source_type.as_mode())?.into_syntax();
     println!("{python_ast:#?}");
     Ok(())
 }

crates/ruff_linter/src/checkers/ast/mod.rs

@@ -32,8 +32,10 @@ use itertools::Itertools;
 use log::debug;
 use ruff_python_ast::{
     self as ast, AnyParameterRef, Comprehension, ElifElseClause, ExceptHandler, Expr, ExprContext,
-    FStringElement, Keyword, MatchCase, Parameter, Parameters, Pattern, Stmt, Suite, UnaryOp,
+    FStringElement, Keyword, MatchCase, ModModule, Parameter, Parameters, Pattern, Stmt, Suite,
+    UnaryOp,
 };
+use ruff_python_parser::Program;
 use ruff_text_size::{Ranged, TextRange, TextSize};
 
 use ruff_diagnostics::{Diagnostic, IsolationLevel};
@@ -174,6 +176,8 @@ impl ExpectedDocstringKind {
 }
 
 pub(crate) struct Checker<'a> {
+    /// The parsed [`Program`].
+    program: &'a Program<ModModule>,
     /// The [`Path`] to the file under analysis.
     path: &'a Path,
     /// The [`Path`] to the package containing the current file.
@@ -223,6 +227,7 @@ pub(crate) struct Checker<'a> {
 impl<'a> Checker<'a> {
     #[allow(clippy::too_many_arguments)]
     pub(crate) fn new(
+        program: &'a Program<ModModule>,
         settings: &'a LinterSettings,
         noqa_line_for: &'a NoqaMapping,
         noqa: flags::Noqa,
@@ -232,12 +237,12 @@ impl<'a> Checker<'a> {
         locator: &'a Locator,
         stylist: &'a Stylist,
         indexer: &'a Indexer,
-        importer: Importer<'a>,
         source_type: PySourceType,
         cell_offsets: Option<&'a CellOffsets>,
         notebook_index: Option<&'a NotebookIndex>,
     ) -> Checker<'a> {
         Checker {
+            program,
             settings,
             noqa_line_for,
             noqa,
@@ -248,7 +253,7 @@ impl<'a> Checker<'a> {
             locator,
             stylist,
             indexer,
-            importer,
+            importer: Importer::new(program.suite(), locator, stylist),
             semantic: SemanticModel::new(&settings.typing_modules, path, module),
             visit: deferred::Visit::default(),
             analyze: deferred::Analyze::default(),
@@ -318,6 +323,11 @@ impl<'a> Checker<'a> {
         }
     }
 
+    /// The [`Program`] for the current file, which contains the tokens, AST, and more.
+    pub(crate) const fn program(&self) -> &'a Program<ModModule> {
+        self.program
+    }
+
     /// The [`Locator`] for the current file, which enables extraction of source code from byte
     /// offsets.
     pub(crate) const fn locator(&self) -> &'a Locator<'a> {
@@ -2321,7 +2331,7 @@ impl<'a> Checker<'a> {
 
 #[allow(clippy::too_many_arguments)]
 pub(crate) fn check_ast(
-    python_ast: &Suite,
+    program: &Program<ModModule>,
     locator: &Locator,
     stylist: &Stylist,
     indexer: &Indexer,
@@ -2351,10 +2361,11 @@ pub(crate) fn check_ast(
         } else {
             ModuleSource::File(path)
         },
-        python_ast,
+        python_ast: program.suite(),
     };
 
     let mut checker = Checker::new(
+        program,
         settings,
         noqa_line_for,
         noqa,
@@ -2364,16 +2375,15 @@ pub(crate) fn check_ast(
         locator,
         stylist,
         indexer,
-        Importer::new(python_ast, locator, stylist),
         source_type,
         cell_offsets,
         notebook_index,
     );
     checker.bind_builtins();
 
     // Iterate over the AST.
-    checker.visit_module(python_ast);
-    checker.visit_body(python_ast);
+    checker.visit_module(program.suite());
+    checker.visit_body(program.suite());
 
     // Visit any deferred syntax nodes. Take care to visit in order, such that we avoid adding
     // new deferred nodes after visiting nodes of that kind. For example, visiting a deferred

crates/ruff_linter/src/fix/edits.rs

@@ -531,8 +531,9 @@ mod tests {
     use test_case::test_case;
 
     use ruff_diagnostics::{Diagnostic, Edit, Fix};
+    use ruff_python_ast::Stmt;
     use ruff_python_codegen::Stylist;
-    use ruff_python_parser::{lexer, parse_expression, parse_suite, Mode};
+    use ruff_python_parser::{lexer, parse, parse_expression, parse_module, Mode};
     use ruff_source_file::Locator;
     use ruff_text_size::{Ranged, TextRange, TextSize};
 
@@ -541,26 +542,29 @@ mod tests {
         add_to_dunder_all, make_redundant_alias, next_stmt_break, trailing_semicolon,
     };
 
+    /// Parse the given source using [`Mode::Module`] and return the first statement.
+    fn parse_first_stmt(source: &str) -> Result<Stmt> {
+        let suite = parse_module(source)?.into_suite();
+        Ok(suite.into_iter().next().unwrap())
+    }
+
     #[test]
     fn find_semicolon() -> Result<()> {
         let contents = "x = 1";
-        let program = parse_suite(contents)?;
-        let stmt = program.first().unwrap();
+        let stmt = parse_first_stmt(contents)?;
         let locator = Locator::new(contents);
         assert_eq!(trailing_semicolon(stmt.end(), &locator), None);
 
         let contents = "x = 1; y = 1";
-        let program = parse_suite(contents)?;
-        let stmt = program.first().unwrap();
+        let stmt = parse_first_stmt(contents)?;
         let locator = Locator::new(contents);
         assert_eq!(
             trailing_semicolon(stmt.end(), &locator),
             Some(TextSize::from(5))
         );
 
         let contents = "x = 1 ; y = 1";
-        let program = parse_suite(contents)?;
-        let stmt = program.first().unwrap();
+        let stmt = parse_first_stmt(contents)?;
         let locator = Locator::new(contents);
         assert_eq!(
             trailing_semicolon(stmt.end(), &locator),
@@ -572,8 +576,7 @@ x = 1 \
   ; y = 1
 "
         .trim();
-        let program = parse_suite(contents)?;
-        let stmt = program.first().unwrap();
+        let stmt = parse_first_stmt(contents)?;
         let locator = Locator::new(contents);
         assert_eq!(
             trailing_semicolon(stmt.end(), &locator),
@@ -614,26 +617,25 @@ x = 1 \
     #[test]
     fn redundant_alias() {
         let contents = "import x, y as y, z as bees";
-        let program = parse_suite(contents).unwrap();
-        let stmt = program.first().unwrap();
+        let stmt = parse_first_stmt(contents)?;
         assert_eq!(
-            make_redundant_alias(["x"].into_iter().map(Cow::from), stmt),
+            make_redundant_alias(["x"].into_iter().map(Cow::from), &stmt),
             vec![Edit::range_replacement(
                 String::from("x as x"),
                 TextRange::new(TextSize::new(7), TextSize::new(8)),
             )],
             "make just one item redundant"
         );
         assert_eq!(
-            make_redundant_alias(vec!["x", "y"].into_iter().map(Cow::from), stmt),
+            make_redundant_alias(vec!["x", "y"].into_iter().map(Cow::from), &stmt),
             vec![Edit::range_replacement(
                 String::from("x as x"),
                 TextRange::new(TextSize::new(7), TextSize::new(8)),
             )],
             "the second item is already a redundant alias"
         );
         assert_eq!(
-            make_redundant_alias(vec!["x", "z"].into_iter().map(Cow::from), stmt),
+            make_redundant_alias(vec!["x", "z"].into_iter().map(Cow::from), &stmt),
             vec![Edit::range_replacement(
                 String::from("x as x"),
                 TextRange::new(TextSize::new(7), TextSize::new(8)),
@@ -661,13 +663,13 @@ x = 1 \
     fn add_to_dunder_all_test(raw: &str, names: &[&str], expect: &str) -> Result<()> {
         let locator = Locator::new(raw);
         let edits = {
-            let expr = parse_expression(raw)?;
+            let expr = parse_expression(raw)?.expr();
             let stylist = Stylist::from_tokens(
                 &lexer::lex(raw, Mode::Expression).collect::<Vec<_>>(),
                 &locator,
             );
             // SUT
-            add_to_dunder_all(names.iter().copied(), &expr, &stylist)
+            add_to_dunder_all(names.iter().copied(), expr, &stylist)
         };
         let diag = {
             use crate::rules::pycodestyle::rules::MissingNewlineAtEndOfFile;

crates/ruff_linter/src/importer/insertion.rs

@@ -321,7 +321,7 @@ mod tests {
 
     use ruff_python_ast::PySourceType;
     use ruff_python_codegen::Stylist;
-    use ruff_python_parser::{parse_suite, Mode};
+    use ruff_python_parser::{parse, parse_module, Mode};
     use ruff_source_file::{LineEnding, Locator};
     use ruff_text_size::TextSize;
 
@@ -330,11 +330,11 @@ mod tests {
     #[test]
     fn start_of_file() -> Result<()> {
         fn insert(contents: &str) -> Result<Insertion> {
-            let program = parse_suite(contents)?;
+            let suite = parse_module(contents)?.into_suite();
             let tokens = ruff_python_parser::tokenize(contents, Mode::Module);
             let locator = Locator::new(contents);
             let stylist = Stylist::from_tokens(&tokens, &locator);
-            Ok(Insertion::start_of_file(&program, &locator, &stylist))
+            Ok(Insertion::start_of_file(&suite, &locator, &stylist))
         }
 
         let contents = "";