omp-py 0.1.0

//! Self-contained embedded `CPython`: a statically linked, free-threaded
//! interpreter (3.14t, GIL off) with the entire stdlib frozen in-memory.
//!
//! [`Engine`] boots `CPython` inside the current process — no Python
//! installation, no disk access. The stdlib ships as marshalled code objects
//! registered wholesale via `PyImport_FrozenModules`, which is
//! per-interpreter machinery: sub-interpreters (`concurrent.interpreters`)
//! get it too, along with the repo-provided Python modules bundled from
//! `crates/py/python` (e.g. `omp_remote`, remote function execution) and
//! the pure-Python packages pinned in `crates/py/requirements.txt`. The
//! one real filesystem path is a site-packages directory for
//! user-installed wheels (native extensions must be dlopen'd from disk).
//!
//! Native Rust modules registered with [`pyo3::append_to_inittab!`] before
//! [`Builder::init`] are importable from Python by name.
//!
//! ```no_run
//! use omp_py::pyo3::{ffi::c_str, prelude::*};
//!
//! let engine = omp_py::Engine::builder().init().expect("boot python");
//! let greet = c_str!("print('hello from embedded python')");
//! engine.attach(|py| py.run(greet, None, None)).unwrap();
//! ```
//!
//! Embedding contract: binaries that should support native wheels must link
//! with `-Wl,-export_dynamic` so extension modules can resolve the `CPython`
//! C-API from the executable at dlopen. This crate's build script applies it
//! to its own binaries; downstream crates need it in their own build script.

use std::{
	env,
	error::Error,
	ffi::CString,
	fmt,
	mem::MaybeUninit,
	os::unix::ffi::OsStrExt,
	path::PathBuf,
	ptr::{null, null_mut},
	sync::atomic::{AtomicBool, Ordering},
};

pub use pyo3;
use pyo3::{ffi, prelude::*};

/// Embedded stdlib: `u32` entry count, then records of `u16` name length
/// (including a trailing NUL), `u8` is-package, `u32` code length,
/// NUL-terminated name, marshalled code object. Uncompressed so the frozen
/// table points straight into the mmap'd binary and unused modules are never
/// paged in. Generated by scripts/fetch-python.sh.
static STDLIB_BLOB: &[u8] = include_bytes!(env!("OMP_STDLIB_BLOB"));

/// Repo-provided Python modules (`crates/py/python`) plus the pure-Python
/// packages pinned in `crates/py/requirements.txt` (e.g. cloudpickle),
/// same format as [`STDLIB_BLOB`]. Packed by build.rs with the vendored
/// interpreter.
static OMP_MODULES_BLOB: &[u8] = include_bytes!(env!("OMP_PY_MODULES_BLOB"));

/// License notices for bundled third-party Python packages
/// (`crates/py/requirements.txt`).
///
/// BSD-style terms require reproducing them in shipped materials. This
/// unreferenced constant is not linked into consumer binaries, so
/// redistributors must surface it explicitly (for example, with a `--licenses`
/// flag) or ship `THIRD-PARTY-NOTICES.txt` beside the artifact.
pub const THIRD_PARTY_LICENSES: &str = include_str!("../THIRD-PARTY-NOTICES.txt");

/// One-shot guard: `CPython` supports a single runtime per process.
static INITIALIZED: AtomicBool = AtomicBool::new(false);

/// Why [`Builder::init`] refused to boot.
///
/// CPython-side boot failures (corrupt frozen data, allocator failure) do
/// not surface here: the interpreter prints its diagnostic and exits the
/// process, per embedding convention.
#[derive(Debug)]
#[non_exhaustive]
pub enum InitError {
	/// The engine was already initialized in this process.
	AlreadyInitialized,
	/// A configured search path contains an interior NUL byte.
	InvalidPath(PathBuf),
}

impl fmt::Display for InitError {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		match self {
			Self::AlreadyInitialized => f.write_str("python engine already initialized"),
			Self::InvalidPath(p) => write!(f, "search path contains NUL byte: {}", p.display()),
		}
	}
}

impl Error for InitError {}

/// Configures and boots the embedded interpreter. Created by
/// [`Engine::builder`].
#[derive(Debug, Default)]
#[must_use = "call .init() to boot the interpreter"]
pub struct Builder {
	site_packages: Option<PathBuf>,
}

impl Builder {
	/// Overrides the site-packages directory (the only real filesystem path
	/// on `sys.path`). Defaults to [`default_site_packages`].
	pub fn site_packages(mut self, dir: impl Into<PathBuf>) -> Self {
		self.site_packages = Some(dir.into());
		self
	}

	/// Boots `CPython`: registers the frozen stdlib, then initializes the
	/// runtime in isolated mode. Callable once per process.
	///
	/// # Errors
	/// [`InitError::AlreadyInitialized`] on repeat calls;
	/// [`InitError::InvalidPath`] if the site-packages path contains NUL.
	pub fn init(self) -> Result<Engine, InitError> {
		if INITIALIZED.swap(true, Ordering::SeqCst) {
			return Err(InitError::AlreadyInitialized);
		}
		let site = self.site_packages.unwrap_or_else(default_site_packages);
		let site_c = CString::new(site.as_os_str().as_bytes())
			.map_err(|_| InitError::InvalidPath(site.clone()))?;
		install_frozen_modules();
		init_python(&site_c);
		Ok(Engine { _priv: () })
	}
}

/// Handle to the booted interpreter; proof that [`Builder::init`] ran.
#[derive(Debug)]
pub struct Engine {
	_priv: (),
}

impl Engine {
	/// Starts configuring an engine.
	pub fn builder() -> Builder {
		Builder::default()
	}

	/// Attaches the current thread to the interpreter and runs `f`.
	/// Equivalent to [`pyo3::Python::attach`], gated on initialization.
	pub fn attach<F, R>(&self, f: F) -> R
	where
		F: for<'py> FnOnce(Python<'py>) -> R,
	{
		Python::attach(f)
	}
}

/// Default wheel directory: `$OMP_PY_SITE` or a home-relative fallback.
///
/// The fallback is `~/.local/share/omp-py/site-packages`. Install into it
/// with any free-threaded 3.14 interpreter, e.g.
/// `uv pip install --python "$(uv python find 3.14t)" --target <dir> numpy`.
pub fn default_site_packages() -> PathBuf {
	env::var_os("OMP_PY_SITE").map_or_else(
		|| {
			env::home_dir()
				.map_or_else(env::temp_dir, |home| home.join(".local/share/omp-py"))
				.join("site-packages")
		},
		PathBuf::from,
	)
}

/// Registers every embedded module (stdlib + repo-provided) as a frozen
/// module. Must run before the interpreter initializes; the table and
/// everything it points at live in the binary's static data.
fn install_frozen_modules() {
	let mut table = Vec::new();
	for blob in [STDLIB_BLOB, OMP_MODULES_BLOB] {
		let count = u32::from_le_bytes(blob[..4].try_into().unwrap()) as usize;
		table.reserve(count + 1);
		let mut rest = &blob[4..];
		for _ in 0..count {
			let name_len = u16::from_le_bytes(rest[..2].try_into().unwrap()) as usize;
			let is_pkg = rest[2];
			let code_len = u32::from_le_bytes(rest[3..7].try_into().unwrap()) as usize;
			let (name, code) = (&rest[7..7 + name_len], &rest[7 + name_len..7 + name_len + code_len]);
			assert_eq!(name[name_len - 1], 0, "blob names must be NUL-terminated");
			table.push(ffi::_frozen {
				name:       name.as_ptr().cast(),
				code:       code.as_ptr(),
				size:       i32::try_from(code_len).unwrap(),
				is_package: i32::from(is_pkg),
			});
			rest = &rest[7 + name_len + code_len..];
		}
	}
	table.push(ffi::_frozen {
		name:       null(),
		code:       null(),
		size:       0,
		is_package: 0,
	});
	// SAFETY: called once, before Py_InitializeFromConfig; the leaked table
	// and the blobs it points into are 'static.
	unsafe {
		ffi::PyImport_FrozenModules = Vec::leak(table).as_ptr();
	}
}

/// Aborts with `CPython`'s diagnostic if `status` signals an init failure.
fn check(status: ffi::PyStatus) {
	// SAFETY: PyStatus is a plain value; both calls are safe on any status
	// and Py_ExitStatusException never returns for failure statuses.
	unsafe {
		if ffi::PyStatus_Exception(status) != 0 {
			ffi::Py_ExitStatusException(status);
		}
	}
}

/// Boots the statically linked interpreter in isolated mode with the
/// site-packages directory as the only module search path.
fn init_python(site_packages: &CString) {
	// SAFETY: standard PyConfig embedding sequence — init, populate, hand to
	// Py_InitializeFromConfig, clear. `config` outlives every borrow of it,
	// and the decoded wide string is freed after CPython copies it.
	unsafe {
		let mut config = MaybeUninit::<ffi::PyConfig>::uninit();
		ffi::PyConfig_InitIsolatedConfig(config.as_mut_ptr());
		let config = config.as_mut_ptr();
		(*config).site_import = 0;
		(*config).write_bytecode = 0;
		(*config).buffered_stdio = 0;
		// CPython gates its own frozen stdlib set (codecs, os, ...) behind
		// this flag outside installed layouts; we have no installed layout.
		(*config).use_frozen_modules = 1;
		check(ffi::PyConfig_SetBytesString(
			config,
			&raw mut (*config).program_name,
			c"omp-py".as_ptr(),
		));
		check(ffi::PyConfig_SetBytesString(
			config,
			&raw mut (*config).stdio_encoding,
			c"utf-8".as_ptr(),
		));
		(*config).module_search_paths_set = 1;
		let wide = ffi::Py_DecodeLocale(site_packages.as_ptr(), null_mut());
		assert!(!wide.is_null(), "failed to decode search path");
		let status = ffi::PyWideStringList_Append(&raw mut (*config).module_search_paths, wide);
		ffi::PyMem_RawFree(wide.cast());
		check(status);
		let status = ffi::Py_InitializeFromConfig(config);
		ffi::PyConfig_Clear(config);
		check(status);
	}
}