Flag to print AST in Jupyter Mode

crates/ruff_dev/src/print_ast.rs

@@ -5,19 +5,26 @@ use std::fs;
 use std::path::PathBuf;
 
 use anyhow::Result;
-use rustpython_parser::ast::Suite;
-use rustpython_parser::Parse;
+use rustpython_parser::{parse, Mode};
 
 #[derive(clap::Args)]
 pub(crate) struct Args {
     /// Python file for which to generate the AST.
     #[arg(required = true)]
     file: PathBuf,
+    /// Run in Jupyter mode i.e., allow line magics.
+    #[arg(long)]
+    jupyter: bool,
 }
 
 pub(crate) fn main(args: &Args) -> Result<()> {
     let contents = fs::read_to_string(&args.file)?;
-    let python_ast = Suite::parse(&contents, &args.file.to_string_lossy())?;
+    let mode = if args.jupyter {
+        Mode::Jupyter
+    } else {
+        Mode::Module
+    };
+    let python_ast = parse(&contents, mode, &args.file.to_string_lossy())?;
     println!("{python_ast:#?}");
     Ok(())
 }