RustPython/src/lib.rs

//! This is the `rustpython` binary. If you're looking to embed RustPython into your application,
//! you're likely looking for the [`rustpython_vm`] crate.
//!
//! You can install `rustpython` with `cargo install rustpython`. If you'd like to inject your
//! own native modules, you can make a binary crate that depends on the `rustpython` crate (and
//! probably [`rustpython_vm`], too), and make a `main.rs` that looks like:
//!
//! ```no_run
//! use rustpython::{InterpreterBuilder, InterpreterBuilderExt};
//! use rustpython_vm::{pymodule, py_freeze};
//!
//! fn main() -> std::process::ExitCode {
//!     let builder = InterpreterBuilder::new().init_stdlib();
//!     // Add a native module using builder.ctx
//!     let my_mod_def = my_mod::module_def(&builder.ctx);
//!     let builder = builder
//!         .add_native_module(my_mod_def)
//!         // Add a frozen module
//!         .add_frozen_modules(py_freeze!(source = "def foo(): pass", module_name = "other_thing"));
//!
//!     rustpython::run(builder)
//! }
//!
//! #[pymodule]
//! mod my_mod {
//!     use rustpython_vm::builtins::PyStrRef;
//!
//!     #[pyfunction]
//!     fn do_thing(x: i32) -> i32 {
//!         x + 1
//!     }
//!
//!     #[pyfunction]
//!     fn other_thing(s: PyStrRef) -> (String, usize) {
//!         let new_string = format!("hello from rust, {}!", s);
//!         let prev_len = s.byte_len();
//!         (new_string, prev_len)
//!     }
//! }
//! ```
//!
//! The binary will have all the standard arguments of a python interpreter (including a REPL!) but
//! it will have your modules loaded into the vm.
//!
//! See [`rustpython_derive`](../rustpython_derive/index.html) crate for documentation on macros used in the example above.

#![allow(clippy::needless_doctest_main)]

#[macro_use]
extern crate log;

#[cfg(feature = "flame-it")]
use vm::Settings;

mod interpreter;
mod settings;
mod shell;

#[cfg(feature = "rustpython-pylib")]
pub use rustpython_pylib as pylib;

use rustpython_vm::{AsObject, PyObjectRef, PyResult, VirtualMachine, scope::Scope};
use std::env;
use std::io::IsTerminal;
use std::process::ExitCode;

pub use interpreter::InterpreterBuilderExt;
pub use rustpython_vm::{self as vm, Interpreter, InterpreterBuilder};
pub use settings::{InstallPipMode, RunMode, parse_opts};
pub use shell::run_shell;

#[cfg(all(
    feature = "ssl",
    not(any(feature = "ssl-rustls", feature = "ssl-openssl")),
    not(target_arch = "wasm32"),
))]
compile_error!(
    "Feature \"ssl\" is now enabled by either \"ssl-rustls\" or \"ssl-openssl\". Do not manually pass \"ssl\" feature. To enable ssl-openssl, use --no-default-features to disable ssl-rustls*"
);

/// The main cli of the `rustpython` interpreter. This function will return `std::process::ExitCode`
/// based on the return code of the python code ran through the cli.
///
/// **Note**: This function provides no way to further initialize the VM after the builder is applied.
/// All VM initialization (adding native modules, init hooks, etc.) must be done through the
/// [`InterpreterBuilder`] parameter before calling this function.
#[must_use]
pub fn run(mut builder: InterpreterBuilder) -> ExitCode {
    env_logger::init();

    // NOTE: This is not a WASI convention. But it will be convenient since POSIX shell always defines it.
    #[cfg(target_os = "wasi")]
    {
        if let Ok(pwd) = env::var("PWD") {
            let _ = env::set_current_dir(pwd);
        };
    }

    let (settings, run_mode) = match parse_opts() {
        Ok(x) => x,
        Err(e) => {
            println!("{e}");
            return ExitCode::FAILURE;
        }
    };

    // don't translate newlines (\r\n <=> \n)
    #[cfg(windows)]
    {
        unsafe extern "C" {
            fn _setmode(fd: i32, flags: i32) -> i32;
        }
        unsafe {
            _setmode(0, libc::O_BINARY);
            _setmode(1, libc::O_BINARY);
            _setmode(2, libc::O_BINARY);
        }
    }

    builder = builder.settings(settings);

    let interp = builder.interpreter();
    let exitcode = cfg_select! {
        feature = "capi" => {{
            let local_vm = interp.enter(|vm| vm.new_thread());
            rustpython_capi::init_main_interpreter(interp);
            let result = local_vm.run(|vm| run_rustpython(vm, run_mode));
            rustpython_capi::get_main_interpreter()
                .take()
                .unwrap()
                .finalize(result.err())
        }},
        _ => interp.run(move |vm| run_rustpython(vm, run_mode)),
    };

    rustpython_vm::host_env::os::exit_code(exitcode)
}

fn get_pip(scope: Scope, vm: &VirtualMachine) -> PyResult<()> {
    let get_getpip = rustpython_vm::py_compile!(
        source = r#"\
__import__("io").TextIOWrapper(
    __import__("urllib.request").request.urlopen("https://bootstrap.pypa.io/get-pip.py")
).read()
"#,
        mode = "eval"
    );
    eprintln!("downloading get-pip.py...");
    let getpip_code = vm.run_code_obj(vm.new_code(get_getpip), vm.new_scope_with_builtins())?;
    let getpip_code: rustpython_vm::builtins::PyStrRef = getpip_code
        .downcast()
        .expect("TextIOWrapper.read() should return str");
    eprintln!("running get-pip.py...");
    vm.run_string(scope, getpip_code.expect_str(), "get-pip.py")?;
    Ok(())
}

fn install_pip(installer: InstallPipMode, scope: Scope, vm: &VirtualMachine) -> PyResult<()> {
    if !cfg!(feature = "ssl") {
        return Err(
            vm.new_system_error("install-pip requires rustpython be build with '--features=ssl'")
        );
    }

    match installer {
        InstallPipMode::Ensurepip => vm.run_module("ensurepip"),
        InstallPipMode::GetPip => get_pip(scope, vm),
    }
}

/// The working directory, when it can be named.
fn current_dir() -> Option<String> {
    env::current_dir().ok()?.into_os_string().into_string().ok()
}

/// `_Py_abspath`: `path` joined onto the working directory. Windows normalizes
/// the result through `GetFullPathNameW`; elsewhere it is left as joined.
fn abs_path(path: &str) -> String {
    if path.is_empty() || path == "." {
        return current_dir().unwrap_or_else(|| path.to_owned());
    }
    if cfg!(windows) {
        return std::path::absolute(path)
            .ok()
            .and_then(|p| p.into_os_string().into_string().ok())
            .unwrap_or_else(|| path.to_owned());
    }
    if std::path::Path::new(path).is_absolute() {
        return path.to_owned();
    }
    // Keep the relative path when the working directory is unavailable.
    current_dir().map_or_else(
        || path.to_owned(),
        |cwd| format!("{cwd}{}{path}", std::path::MAIN_SEPARATOR),
    )
}

/// `_PyPathConfig_ComputeSysPath0` for a script argument: the directory holding
/// the script, with symlinks resolved.
fn script_sys_path0(argv0: &str) -> String {
    // `realpath()`, which Windows has no counterpart for; there the path is
    // only made absolute.
    let resolved = if cfg!(windows) {
        Some(abs_path(argv0))
    } else {
        std::fs::canonicalize(argv0)
            .ok()
            .and_then(|p| p.into_os_string().into_string().ok())
    };
    let path0 = resolved.as_deref().unwrap_or(argv0);
    let Some(sep) = path0.rfind(std::path::is_separator) else {
        return String::new();
    };
    // Drop the trailing separator unless it is all that is left of a root.
    let end = if sep == 0 || (cfg!(windows) && path0[..sep].ends_with(':')) {
        sep + 1
    } else {
        sep
    };
    path0[..end].to_owned()
}

// pymain_run_file_obj in Modules/main.c
fn run_file(vm: &VirtualMachine, scope: Scope, argv0: &str) -> PyResult<()> {
    // config_run_filename_abspath: __file__ and co_filename are absolute even
    // when the script was named relative to the working directory.
    let path = &abs_path(argv0);

    // Check if path is a package/directory with __main__.py
    if let Some(_importer) = get_importer(path, vm)? {
        vm.insert_sys_path(vm.new_pyobj(path.clone()))?;
        let runpy = vm.import("runpy", 0)?;
        let run_module_as_main = runpy.get_attr("_run_module_as_main", vm)?;
        run_module_as_main.call((vm::identifier!(vm, __main__).to_owned(), false), vm)?;
        return Ok(());
    }

    // Add script directory to sys.path[0]
    if !vm.state.config.settings.safe_path {
        vm.insert_sys_path(vm.new_pyobj(script_sys_path0(argv0)))?;
    }

    cfg_select! {
        feature = "host_env" => match rustpython_vm::host_env::fs::metadata(path) {
            Ok(_) => vm.run_any_file(scope, path),
            Err(err) => cant_open_file(vm, path, &err),
        },
        _ => {
            // In sandbox mode, the binary reads the file and feeds source to the VM.
            // The VM itself has no filesystem access.
            let path = if path.is_empty() { "???" } else { path };
            match std::fs::read_to_string(path) {
                Ok(source) => vm.run_string(scope, &source, path).map(drop),
                Err(err) => cant_open_file(vm, path, &err),
            }
        }
    }
}

fn cant_open_file(vm: &VirtualMachine, path: &str, err: &std::io::Error) -> PyResult<()> {
    let program = &vm.state.config.paths.executable;
    let filename_repr = vm.ctx.new_str(path).as_object().repr(vm)?;
    let errno = err.raw_os_error().unwrap_or(2);
    eprintln!("{program}: can't open file {filename_repr}: [Errno {errno}] {err}");
    Err(vm.new_system_exit(vec![vm.ctx.new_int(2).into()].into()))
}

fn get_importer(path: &str, vm: &VirtualMachine) -> PyResult<Option<PyObjectRef>> {
    use rustpython_vm::builtins::PyDictRef;
    use rustpython_vm::convert::TryFromObject;

    let path_importer_cache = vm.sys_module.get_attr("path_importer_cache", vm)?;
    let path_importer_cache = PyDictRef::try_from_object(vm, path_importer_cache)?;
    if let Some(importer) = path_importer_cache.get_item_opt(path, vm)? {
        return Ok(Some(importer));
    }
    let path_obj = vm.ctx.new_str(path);
    let path_hooks = vm.sys_module.get_attr("path_hooks", vm)?;
    let mut importer = None;
    let path_hooks: Vec<PyObjectRef> = path_hooks.try_into_value(vm)?;
    for path_hook in path_hooks {
        match path_hook.call((path_obj.clone(),), vm) {
            Ok(imp) => {
                importer = Some(imp);
                break;
            }
            Err(e) if e.fast_isinstance(vm.ctx.exceptions.import_error) => continue,
            Err(e) => return Err(e),
        }
    }
    Ok(if let Some(imp) = importer {
        let imp = path_importer_cache.get_or_insert(vm, path_obj.as_object(), || imp.clone())?;
        Some(imp)
    } else {
        None
    })
}

// pymain_run_python
fn run_rustpython(vm: &VirtualMachine, run_mode: RunMode) -> PyResult<()> {
    #[cfg(feature = "flame-it")]
    let main_guard = flame::start_guard("RustPython main");

    let scope = vm.new_scope_with_main()?;

    // Initialize warnings module to process sys.warnoptions
    // _PyWarnings_Init()
    if vm.import("warnings", 0).is_err() {
        warn!("Failed to import warnings module");
    }

    // Import site first, before setting sys.path[0]
    // This matches CPython's behavior where site.removeduppaths() runs
    // before sys.path[0] is set, preventing '' from being converted to cwd
    if vm.state.config.settings.import_site && vm.import("site", 0).is_err() {
        warn!(
            "Failed to import site, consider adding the Lib directory to your RUSTPYTHONPATH \
             environment variable",
        );
    }

    // _PyPathConfig_ComputeSysPath0 - set sys.path[0] after site import
    if !vm.state.config.settings.safe_path {
        let path0: Option<String> = match &run_mode {
            RunMode::Command(_) => Some(String::new()),
            RunMode::Module(_) => env::current_dir()
                .ok()
                .and_then(|p| p.to_str().map(|s| s.to_owned())),
            RunMode::Script(_) | RunMode::InstallPip(_) => None, // handled by run_script
            RunMode::Repl => Some(String::new()),
        };

        if let Some(path) = path0 {
            vm.insert_sys_path(vm.new_pyobj(path))?;
        }
    }

    // Enable faulthandler if -X faulthandler, PYTHONFAULTHANDLER or -X dev is set
    // _PyFaulthandler_Init()
    if vm.state.config.settings.faulthandler {
        let _ = vm.run_simple_string("import faulthandler; faulthandler.enable()");
    }

    let is_repl = matches!(run_mode, RunMode::Repl);
    if !vm.state.config.settings.quiet
        && (vm.state.config.settings.verbose > 0 || (is_repl && std::io::stdin().is_terminal()))
    {
        eprintln!(
            "Welcome to the magnificent Rust Python {} interpreter \u{1f631} \u{1f596}",
            env!("CARGO_PKG_VERSION")
        );
        eprintln!(
            "RustPython {}.{}.{}",
            vm::version::MAJOR,
            vm::version::MINOR,
            vm::version::MICRO,
        );

        eprintln!("Type \"help\", \"copyright\", \"credits\" or \"license\" for more information.");
    }
    let res = match run_mode {
        RunMode::Command(command) => {
            debug!("Running command {command}");
            vm.run_string(scope.clone(), &command, "<string>").map(drop)
        }
        RunMode::Module(module) => {
            debug!("Running module {module}");
            vm.run_module(&module)
        }
        RunMode::InstallPip(installer) => install_pip(installer, scope.clone(), vm),
        RunMode::Script(script_path) => {
            // pymain_run_file_obj
            debug!("Running script {script_path}");
            run_file(vm, scope.clone(), &script_path)
        }
        RunMode::Repl => Ok(()),
    };
    let result = if is_repl || vm.state.config.settings.inspect {
        if std::io::stdin().is_terminal() {
            warn_if_pyrepl_unavailable(vm);
        }
        shell::run_shell(vm, scope)
    } else {
        res
    };

    #[cfg(feature = "flame-it")]
    {
        main_guard.end();
        if let Err(e) = write_profile(&vm.state.as_ref().config.settings) {
            error!("Error writing profile information: {}", e);
        }
    }

    result
}

/// When the fancy REPL cannot start (for example `TERM=dumb`), print the same
/// warning `_pyrepl.main` would have printed, then keep the rustyline shell.
fn warn_if_pyrepl_unavailable(vm: &VirtualMachine) {
    let _ = vm.run_simple_string(
        r"
import os, sys
if not os.getenv('PYTHON_BASIC_REPL'):
    try:
        from _pyrepl.main import CAN_USE_PYREPL, FAIL_REASON
        if not CAN_USE_PYREPL and FAIL_REASON:
            print(FAIL_REASON, file=sys.stderr)
    except Exception:
        pass
",
    );
}

#[cfg(feature = "flame-it")]
fn write_profile(settings: &Settings) -> Result<(), Box<dyn core::error::Error>> {
    use std::{fs, io};

    enum ProfileFormat {
        Html,
        Text,
        SpeedScope,
    }
    let profile_output = settings.profile_output.as_deref();
    let profile_format = match settings.profile_format.as_deref() {
        Some("html") => ProfileFormat::Html,
        Some("text") => ProfileFormat::Text,
        None if profile_output == Some("-".as_ref()) => ProfileFormat::Text,
        // spell-checker:ignore speedscope
        Some("speedscope") | None => ProfileFormat::SpeedScope,
        Some(other) => {
            error!("Unknown profile format {}", other);
            // TODO: Need to change to ExitCode or Termination
            std::process::exit(1);
        }
    };

    let profile_output = profile_output.unwrap_or_else(|| match profile_format {
        ProfileFormat::Html => "flame-graph.html".as_ref(),
        ProfileFormat::Text => "flame.txt".as_ref(),
        ProfileFormat::SpeedScope => "flamescope.json".as_ref(),
    });

    let profile_output: Box<dyn io::Write> = if profile_output == "-" {
        Box::new(io::stdout())
    } else {
        Box::new(fs::File::create(profile_output)?)
    };

    let profile_output = io::BufWriter::new(profile_output);

    match profile_format {
        ProfileFormat::Html => flame::dump_html(profile_output)?,
        ProfileFormat::Text => flame::dump_text_to_writer(profile_output)?,
        ProfileFormat::SpeedScope => flamescope::dump(profile_output)?,
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use rustpython_vm::Interpreter;

    fn interpreter() -> Interpreter {
        InterpreterBuilder::new().init_stdlib().interpreter()
    }

    #[test]
    fn run_script() {
        interpreter().enter(|vm| {
            vm.unwrap_pyresult((|| {
                let scope = vm.new_scope_with_main()?;
                // test file run
                run_file(vm, scope, "extra_tests/snippets/dir_main/__main__.py")?;

                #[cfg(feature = "host_env")]
                {
                    let scope = vm.new_scope_with_main()?;
                    // test module run (directory with __main__.py)
                    run_file(vm, scope, "extra_tests/snippets/dir_main")?;
                }

                Ok(())
            })());
        })
    }
}