1#![allow(clippy::needless_doctest_main)]
39
40#[macro_use]
41extern crate clap;
42extern crate env_logger;
43#[macro_use]
44extern crate log;
45
46#[cfg(feature = "flame-it")]
47use vm::Settings;
48
49mod interpreter;
50mod settings;
51mod shell;
52
53use atty::Stream;
54use rustpython_vm::{scope::Scope, PyResult, VirtualMachine};
55use std::{env, process::ExitCode};
56
57pub use interpreter::InterpreterConfig;
58pub use rustpython_vm as vm;
59pub use settings::{opts_with_clap, RunMode};
60
61pub fn run(init: impl FnOnce(&mut VirtualMachine) + 'static) -> ExitCode {
64 env_logger::init();
65
66 #[cfg(target_os = "wasi")]
68 {
69 if let Ok(pwd) = env::var("PWD") {
70 let _ = env::set_current_dir(pwd);
71 };
72 }
73
74 let (settings, run_mode) = opts_with_clap();
75
76 let quiet_var = settings.quiet || !atty::is(Stream::Stdin);
78
79 #[cfg(windows)]
81 {
82 extern "C" {
83 fn _setmode(fd: i32, flags: i32) -> i32;
84 }
85 unsafe {
86 _setmode(0, libc::O_BINARY);
87 _setmode(1, libc::O_BINARY);
88 _setmode(2, libc::O_BINARY);
89 }
90 }
91
92 let mut config = InterpreterConfig::new().settings(settings);
93 #[cfg(feature = "stdlib")]
94 {
95 config = config.init_stdlib();
96 }
97 config = config.init_hook(Box::new(init));
98
99 let interp = config.interpreter();
100 let exitcode = interp.run(move |vm| run_rustpython(vm, run_mode, quiet_var));
101
102 ExitCode::from(exitcode)
103}
104
105fn setup_main_module(vm: &VirtualMachine) -> PyResult<Scope> {
106 let scope = vm.new_scope_with_builtins();
107 let main_module = vm.new_module("__main__", scope.globals.clone(), None);
108 main_module
109 .dict()
110 .set_item("__annotations__", vm.ctx.new_dict().into(), vm)
111 .expect("Failed to initialize __main__.__annotations__");
112
113 vm.sys_module
114 .get_attr("modules", vm)?
115 .set_item("__main__", main_module.into(), vm)?;
116
117 Ok(scope)
118}
119
120#[cfg(feature = "ssl")]
121fn get_pip(scope: Scope, vm: &VirtualMachine) -> PyResult<()> {
122 let get_getpip = rustpython_vm::py_compile!(
123 source = r#"\
124__import__("io").TextIOWrapper(
125 __import__("urllib.request").request.urlopen("https://bootstrap.pypa.io/get-pip.py")
126).read()
127"#,
128 mode = "eval"
129 );
130 eprintln!("downloading get-pip.py...");
131 let getpip_code = vm.run_code_obj(vm.ctx.new_code(get_getpip), scope.clone())?;
132 let getpip_code: rustpython_vm::builtins::PyStrRef = getpip_code
133 .downcast()
134 .expect("TextIOWrapper.read() should return str");
135 eprintln!("running get-pip.py...");
136 vm.run_code_string(scope, getpip_code.as_str(), "get-pip.py".to_owned())?;
137 Ok(())
138}
139
140#[cfg(feature = "ssl")]
141fn ensurepip(_: Scope, vm: &VirtualMachine) -> PyResult<()> {
142 vm.run_module("ensurepip")
143}
144
145fn install_pip(_installer: &str, _scope: Scope, vm: &VirtualMachine) -> PyResult<()> {
146 #[cfg(feature = "ssl")]
147 {
148 match _installer {
149 "ensurepip" => ensurepip(_scope, vm),
150 "get-pip" => get_pip(_scope, vm),
151 _ => unreachable!(),
152 }
153 }
154
155 #[cfg(not(feature = "ssl"))]
156 Err(vm.new_exception_msg(
157 vm.ctx.exceptions.system_error.to_owned(),
158 "install-pip requires rustpython be build with '--features=ssl'".to_owned(),
159 ))
160}
161
162fn run_rustpython(vm: &VirtualMachine, run_mode: RunMode, quiet: bool) -> PyResult<()> {
163 #[cfg(feature = "flame-it")]
164 let main_guard = flame::start_guard("RustPython main");
165
166 let scope = setup_main_module(vm)?;
167
168 if !vm.state.settings.safe_path {
169 vm.run_code_string(
172 vm.new_scope_with_builtins(),
173 "import sys; sys.path.insert(0, '')",
174 "<embedded>".to_owned(),
175 )?;
176 }
177
178 let site_result = vm.import("site", 0);
179 if site_result.is_err() {
180 warn!(
181 "Failed to import site, consider adding the Lib directory to your RUSTPYTHONPATH \
182 environment variable",
183 );
184 }
185
186 match run_mode {
187 RunMode::Command(command) => {
188 debug!("Running command {}", command);
189 vm.run_code_string(scope, &command, "<stdin>".to_owned())?;
190 }
191 RunMode::Module(module) => {
192 debug!("Running module {}", module);
193 vm.run_module(&module)?;
194 }
195 RunMode::InstallPip(installer) => {
196 install_pip(&installer, scope, vm)?;
197 }
198 RunMode::ScriptInteractive(script, interactive) => {
199 if let Some(script) = script {
200 debug!("Running script {}", &script);
201 vm.run_script(scope.clone(), &script)?;
202 } else if !quiet {
203 println!(
204 "Welcome to the magnificent Rust Python {} interpreter \u{1f631} \u{1f596}",
205 crate_version!()
206 );
207 }
208 if interactive {
209 shell::run_shell(vm, scope)?;
210 }
211 }
212 }
213 #[cfg(feature = "flame-it")]
214 {
215 main_guard.end();
216 if let Err(e) = write_profile(&vm.state.as_ref().settings) {
217 error!("Error writing profile information: {}", e);
218 }
219 }
220 Ok(())
221}
222
223#[cfg(feature = "flame-it")]
224fn write_profile(settings: &Settings) -> Result<(), Box<dyn std::error::Error>> {
225 use std::{fs, io};
226
227 enum ProfileFormat {
228 Html,
229 Text,
230 Speedscope,
231 }
232 let profile_output = settings.profile_output.as_deref();
233 let profile_format = match settings.profile_format.as_deref() {
234 Some("html") => ProfileFormat::Html,
235 Some("text") => ProfileFormat::Text,
236 None if profile_output == Some("-".as_ref()) => ProfileFormat::Text,
237 Some("speedscope") | None => ProfileFormat::Speedscope,
238 Some(other) => {
239 error!("Unknown profile format {}", other);
240 std::process::exit(1);
242 }
243 };
244
245 let profile_output = profile_output.unwrap_or_else(|| match profile_format {
246 ProfileFormat::Html => "flame-graph.html".as_ref(),
247 ProfileFormat::Text => "flame.txt".as_ref(),
248 ProfileFormat::Speedscope => "flamescope.json".as_ref(),
249 });
250
251 let profile_output: Box<dyn io::Write> = if profile_output == "-" {
252 Box::new(io::stdout())
253 } else {
254 Box::new(fs::File::create(profile_output)?)
255 };
256
257 let profile_output = io::BufWriter::new(profile_output);
258
259 match profile_format {
260 ProfileFormat::Html => flame::dump_html(profile_output)?,
261 ProfileFormat::Text => flame::dump_text_to_writer(profile_output)?,
262 ProfileFormat::Speedscope => flamescope::dump(profile_output)?,
263 }
264
265 Ok(())
266}
267
268#[cfg(test)]
269mod tests {
270 use super::*;
271 use rustpython_vm::Interpreter;
272
273 fn interpreter() -> Interpreter {
274 InterpreterConfig::new().init_stdlib().interpreter()
275 }
276
277 #[test]
278 fn test_run_script() {
279 interpreter().enter(|vm| {
280 vm.unwrap_pyresult((|| {
281 let scope = setup_main_module(vm)?;
282 vm.run_script(scope, "extra_tests/snippets/dir_main/__main__.py")?;
284
285 let scope = setup_main_module(vm)?;
286 vm.run_script(scope, "extra_tests/snippets/dir_main")?;
288
289 Ok(())
290 })());
291 })
292 }
293}