pyforge 0.3.0

//! Implementation details of `#[pymodule]` which need to be accessible from proc-macro generated code.

use std::{
    cell::UnsafeCell,
    ffi::CStr,
    marker::PhantomData,
    os::raw::{c_int, c_void},
};

#[cfg(all(
    Py_3_9,
    not(all(windows, Py_LIMITED_API, not(Py_3_10))),
))]
use std::sync::atomic::Ordering;

#[cfg(all(
    Py_3_9,
    not(all(windows, Py_LIMITED_API, not(Py_3_10))),
    not(target_has_atomic = "64"),
))]
use portable_atomic::AtomicI64;
#[cfg(all(
    Py_3_9,
    not(all(windows, Py_LIMITED_API, not(Py_3_10))),
    target_has_atomic = "64",
))]
use std::sync::atomic::AtomicI64;

use crate::exceptions::PyImportError;
use crate::prelude::PyTypeMethods;
use crate::{
    ffi,
    impl_::pyfunction::PyFunctionDef,
    sync::PyOnceLock,
    types::{any::PyAnyMethods, dict::PyDictMethods, PyDict, PyModule, PyModuleMethods},
    Bound, Py, PyAny, PyClass, PyResult, PyTypeInfo, Python,
};
use crate::{ffi_ptr_ext::FfiPtrExt, PyErr};

/// `Sync` wrapper of `ffi::PyModuleDef`.
pub struct ModuleDef {
    // wrapped in UnsafeCell so that Rust compiler treats this as interior mutability
    ffi_def: UnsafeCell<ffi::PyModuleDef>,
    #[cfg(Py_3_15)]
    name: &'static CStr,
    #[cfg(Py_3_15)]
    doc: &'static CStr,
    slots: &'static PyModuleSlots,
    /// Interpreter ID where module was initialized.
    #[cfg(all(
        Py_3_9,
        not(all(windows, Py_LIMITED_API, not(Py_3_10)))
    ))]
    interpreter: AtomicI64,
    /// Initialized module object, cached to avoid reinitialization.
    module: PyOnceLock<Py<PyModule>>,
}

unsafe impl Sync for ModuleDef {}

impl ModuleDef {
    /// Make new module definition with given module name.
    pub const fn new(
        name: &'static CStr,
        doc: &'static CStr,
        slots: &'static PyModuleSlots,
    ) -> Self {
        // This is only used in PyForge for append_to_inittab on Python 3.15 and newer.
        // There could also be other tools that need the legacy init hook.
        // Opaque PyObject builds won't be able to use this.
        #[allow(clippy::declare_interior_mutable_const)]
        const INIT: ffi::PyModuleDef = ffi::PyModuleDef {
            m_base: ffi::PyModuleDef_HEAD_INIT,
            m_name: std::ptr::null(),
            m_doc: std::ptr::null(),
            m_size: 0,
            m_methods: std::ptr::null_mut(),
            m_slots: std::ptr::null_mut(),
            m_traverse: None,
            m_clear: None,
            m_free: None,
        };

        let ffi_def = UnsafeCell::new(ffi::PyModuleDef {
            m_name: name.as_ptr(),
            m_doc: doc.as_ptr(),
            // TODO: would be slightly nicer to use `[T]::as_mut_ptr()` here,
            // but that requires mut ptr deref on MSRV.
            m_slots: slots.0.get() as _,
            ..INIT
        });

        ModuleDef {
            ffi_def,
            #[cfg(Py_3_15)]
            name,
            #[cfg(Py_3_15)]
            doc,
            slots,
            // -1 is never expected to be a valid interpreter ID
            #[cfg(all(
                Py_3_9,
                not(all(windows, Py_LIMITED_API, not(Py_3_10)))
            ))]
            interpreter: AtomicI64::new(-1),
            module: PyOnceLock::new(),
        }
    }

    pub fn init_multi_phase(&'static self) -> *mut ffi::PyObject {
        unsafe { ffi::PyModuleDef_Init(self.ffi_def.get()) }
    }

    /// Builds a module object directly. Used for [`#[pymodule]`][crate::pymodule] submodules.
    pub fn make_module(&'static self, py: Python<'_>) -> PyResult<Py<PyModule>> {
        // Check the interpreter ID has not changed, since we currently have no way to guarantee
        // that static data is not reused across interpreters.
        //
        // TODO: it should be possible to use the Py_mod_multiple_interpreters slot on sufficiently
        // new Python versions to remove the need for this custom logic
        {
            // PyInterpreterState_Get is only available on 3.9 and later, but is missing
            // from python3.dll for Windows stable API on 3.9
            #[cfg(all(Py_3_9, not(all(windows, Py_LIMITED_API, not(Py_3_10)))))]
            {
                let current_interpreter =
                    unsafe { ffi::PyInterpreterState_GetID(ffi::PyInterpreterState_Get()) };
                crate::err::error_on_minusone(py, current_interpreter)?;
                if let Err(initialized_interpreter) = self.interpreter.compare_exchange(
                    -1,
                    current_interpreter,
                    Ordering::SeqCst,
                    Ordering::SeqCst,
                ) {
                    if initialized_interpreter != current_interpreter {
                        return Err(PyImportError::new_err(
                            "PyForge modules do not yet support subinterpreters, see https://github.com/PyForge/pyo3/issues/576",
                        ));
                    }
                }
            }
            #[cfg(not(all(Py_3_9, not(all(windows, Py_LIMITED_API, not(Py_3_10))))))]
            {
                // CPython before 3.9 does not have APIs to check the interpreter ID, so best that can be
                // done to guard against subinterpreters is fail if the module is initialized twice
                if self.module.get(py).is_some() {
                    return Err(PyImportError::new_err(
                        "PyForge modules compiled for CPython 3.8 or older may only be initialized once per interpreter process"
                    ));
                }
            }
        }

        // Make a dummy spec, needs a `name` attribute and that seems to be sufficient
        // for the loader system

        static SIMPLE_NAMESPACE: PyOnceLock<Py<PyAny>> = PyOnceLock::new();
        let simple_ns = SIMPLE_NAMESPACE.import(py, "types", "SimpleNamespace")?;

        #[cfg(not(Py_3_15))]
        {
            let ffi_def = self.ffi_def.get();

            let m_name = unsafe { CStr::from_ptr((*ffi_def).m_name) };
            let name = m_name
                .to_str()
                .map_err(|e| {
                    crate::exceptions::PyUnicodeDecodeError::new_err_from_utf8(
                        py,
                        m_name.to_bytes(),
                        e,
                    )
                })?
                .to_string();
            let kwargs = PyDict::new(py);
            kwargs.set_item("name", name)?;
            let spec = simple_ns.call((), Some(&kwargs))?;

            self.module
                .get_or_try_init(py, || {
                    let def = self.ffi_def.get();
                    let module = unsafe {
                        ffi::PyModule_FromDefAndSpec(def, spec.as_ptr()).assume_owned_or_err(py)?
                    }
                    .cast_into()?;
                    if unsafe { ffi::PyModule_ExecDef(module.as_ptr(), def) } != 0 {
                        return Err(PyErr::fetch(py));
                    }
                    Ok(module.unbind())
                })
                .map(|py_module| py_module.clone_ref(py))
        }

        #[cfg(Py_3_15)]
        {
            let name = self.name;
            let doc = self.doc;
            let kwargs = PyDict::new(py);
            kwargs.set_item("name", name)?;
            let spec = simple_ns.call((), Some(&kwargs))?;

            self.module
                .get_or_try_init(py, || {
                    let slots = self.get_slots();
                    let module = unsafe {
                        ffi::PyModule_FromSlotsAndSpec(slots, spec.as_ptr())
                            .assume_owned_or_err(py)?
                    }
                    .cast_into()?;
                    if unsafe { ffi::PyModule_SetDocString(module.as_ptr(), doc.as_ptr()) } != 0 {
                        return Err(PyErr::fetch(py));
                    }
                    if unsafe { ffi::PyModule_Exec(module.as_ptr()) } != 0 {
                        return Err(PyErr::fetch(py));
                    }
                    Ok(module.unbind())
                })
                .map(|py_module| py_module.clone_ref(py))
        }
    }
    pub fn get_slots(&'static self) -> *mut ffi::PyModuleDef_Slot {
        self.slots.0.get() as *mut ffi::PyModuleDef_Slot
    }
}

/// Type of the exec slot used to initialise module contents
pub type ModuleExecSlot = unsafe extern "C" fn(*mut ffi::PyObject) -> c_int;

const MAX_SLOTS: usize =
    // Py_mod_exec
    1 +
    // Py_mod_gil
    cfg!(Py_3_13) as usize +
    // Py_mod_name, Py_mod_doc, and Py_mod_abi
    3 * (cfg!(Py_3_15) as usize);
const MAX_SLOTS_WITH_TRAILING_NULL: usize = MAX_SLOTS + 1;

/// Builder to create `PyModuleSlots`. The size of the number of slots desired must
/// be known up front, and N needs to be at least one greater than the number of
/// actual slots pushed due to the need to have a zeroed element on the end.
pub struct PyModuleSlotsBuilder {
    // values (initially all zeroed)
    values: [ffi::PyModuleDef_Slot; MAX_SLOTS_WITH_TRAILING_NULL],
    // current length
    len: usize,
}

// note that macros cannot use conditional compilation,
// so all implementations below must be available in all
// Python versions
// By handling it here we can avoid conditional
// compilation within the macros; they can always emit
// e.g. a `.with_gil_used()` call.
impl PyModuleSlotsBuilder {
    #[allow(clippy::new_without_default)]
    pub const fn new() -> Self {
        Self {
            values: [unsafe { std::mem::zeroed() }; MAX_SLOTS_WITH_TRAILING_NULL],
            len: 0,
        }
    }

    pub const fn with_mod_exec(self, exec: ModuleExecSlot) -> Self {
        self.push(ffi::Py_mod_exec, exec as *mut c_void)
    }

    pub const fn with_gil_used(self, gil_used: bool) -> Self {
        #[cfg(Py_3_13)]
        {
            self.push(
                ffi::Py_mod_gil,
                if gil_used {
                    ffi::Py_MOD_GIL_USED
                } else {
                    ffi::Py_MOD_GIL_NOT_USED
                },
            )
        }

        #[cfg(not(Py_3_13))]
        {
            // Silence unused variable warning
            let _ = gil_used;
            self
        }
    }

    pub const fn with_name(self, name: &'static CStr) -> Self {
        #[cfg(Py_3_15)]
        {
            self.push(ffi::Py_mod_name, name.as_ptr() as *mut c_void)
        }

        #[cfg(not(Py_3_15))]
        {
            // Silence unused variable warning
            let _ = name;
            self
        }
    }

    pub const fn with_abi_info(self) -> Self {
        #[cfg(Py_3_15)]
        {
            ffi::PyABIInfo_VAR!(ABI_INFO);
            self.push(ffi::Py_mod_abi, std::ptr::addr_of_mut!(ABI_INFO).cast())
        }

        #[cfg(not(Py_3_15))]
        {
            self
        }
    }

    pub const fn with_doc(self, doc: &'static CStr) -> Self {
        #[cfg(Py_3_15)]
        {
            self.push(ffi::Py_mod_doc, doc.as_ptr() as *mut c_void)
        }

        #[cfg(not(Py_3_15))]
        {
            // Silence unused variable warning
            let _ = doc;
            self
        }
    }

    pub const fn build(self) -> PyModuleSlots {
        PyModuleSlots(UnsafeCell::new(self.values))
    }

    const fn push(mut self, slot: c_int, value: *mut c_void) -> Self {
        // Required to guarantee there's still a zeroed element
        // at the end
        assert!(
            self.len < MAX_SLOTS,
            "Cannot add more than MAX_SLOTS slots to a PyModuleSlots",
        );
        self.values[self.len] = ffi::PyModuleDef_Slot { slot, value };
        self.len += 1;
        self
    }
}

/// Wrapper to safely store module slots, to be used in a `ModuleDef`.
pub struct PyModuleSlots(UnsafeCell<[ffi::PyModuleDef_Slot; MAX_SLOTS_WITH_TRAILING_NULL]>);

// It might be possible to avoid this with SyncUnsafeCell in the future
//
// SAFETY: the inner values are only accessed within a `ModuleDef`,
// which only uses them to build the `ffi::ModuleDef`.
unsafe impl Sync for PyModuleSlots {}

/// Trait to add an element (class, function...) to a module.
///
/// Currently only implemented for classes.
pub trait PyAddToModule: crate::sealed::Sealed {
    fn add_to_module(&'static self, module: &Bound<'_, PyModule>) -> PyResult<()>;
}

/// For adding native types (non-pyclass) to a module.
pub struct AddTypeToModule<T>(PhantomData<T>);

impl<T> AddTypeToModule<T> {
    #[allow(clippy::new_without_default)]
    pub const fn new() -> Self {
        AddTypeToModule(PhantomData)
    }
}

impl<T: PyTypeInfo> PyAddToModule for AddTypeToModule<T> {
    fn add_to_module(&'static self, module: &Bound<'_, PyModule>) -> PyResult<()> {
        let object = T::type_object(module.py());
        module.add(object.name()?, object)
    }
}

/// For adding a class to a module.
pub struct AddClassToModule<T>(PhantomData<T>);

impl<T> AddClassToModule<T> {
    #[allow(clippy::new_without_default)]
    pub const fn new() -> Self {
        AddClassToModule(PhantomData)
    }
}

impl<T: PyClass> PyAddToModule for AddClassToModule<T> {
    fn add_to_module(&'static self, module: &Bound<'_, PyModule>) -> PyResult<()> {
        module.add_class::<T>()
    }
}

/// For adding a function to a module.
impl PyAddToModule for PyFunctionDef {
    fn add_to_module(&'static self, module: &Bound<'_, PyModule>) -> PyResult<()> {
        // safety: self is static
        module.add_function(self.create_py_c_function(module.py(), Some(module))?)
    }
}

/// For adding a module to a module.
impl PyAddToModule for ModuleDef {
    fn add_to_module(&'static self, module: &Bound<'_, PyModule>) -> PyResult<()> {
        module.add_submodule(self.make_module(module.py())?.bind(module.py()))
    }
}

#[cfg(test)]
mod tests {
    use std::{borrow::Cow, ffi::CStr, os::raw::c_int};

    use crate::{
        ffi,
        impl_::{
            pymodule::{PyModuleSlots, PyModuleSlotsBuilder},
            trampoline,
        },
        types::{any::PyAnyMethods, module::PyModuleMethods},
        Python,
    };

    use super::{ModuleDef, MAX_SLOTS};

    unsafe extern "C" fn module_exec(_module: *mut ffi::PyObject) -> c_int {
        0
    }

    #[test]
    fn module_init() {
        unsafe extern "C" fn module_exec(module: *mut ffi::PyObject) -> c_int {
            unsafe {
                trampoline::module_exec(module, |m| {
                    m.add("SOME_CONSTANT", 42)?;
                    Ok(())
                })
            }
        }

        static NAME: &CStr = c"test_module";
        static DOC: &CStr = c"some doc";

        static SLOTS: PyModuleSlots = PyModuleSlotsBuilder::new()
            .with_mod_exec(module_exec)
            .with_gil_used(false)
            .with_abi_info()
            .with_name(NAME)
            .with_doc(DOC)
            .build();

        static MODULE_DEF: ModuleDef = ModuleDef::new(NAME, DOC, &SLOTS);

        Python::attach(|py| {
            let module = MODULE_DEF.make_module(py).unwrap().into_bound(py);
            assert_eq!(
                module
                    .getattr("__name__")
                    .unwrap()
                    .extract::<Cow<'_, str>>()
                    .unwrap(),
                "test_module",
            );
            assert_eq!(
                module
                    .getattr("__doc__")
                    .unwrap()
                    .extract::<Cow<'_, str>>()
                    .unwrap(),
                "some doc",
            );
            assert_eq!(
                module
                    .getattr("SOME_CONSTANT")
                    .unwrap()
                    .extract::<u8>()
                    .unwrap(),
                42,
            );
        })
    }

    #[test]
    fn module_def_new() {
        // To get coverage for ModuleDef::new() need to create a non-static ModuleDef, however init
        // etc require static ModuleDef, so this test needs to be separated out.
        static NAME: &CStr = c"test_module";
        static DOC: &CStr = c"some doc";

        static SLOTS: PyModuleSlots = PyModuleSlotsBuilder::new().build();

        let module_def: ModuleDef = ModuleDef::new(NAME, DOC, &SLOTS);

        unsafe {
            assert_eq!((*module_def.ffi_def.get()).m_slots, SLOTS.0.get().cast());
        }
        #[cfg(Py_3_15)]
        {
            assert_eq!(module_def.name, NAME);
            assert_eq!(module_def.doc, DOC);
        }
        assert_eq!(module_def.slots.0.get(), SLOTS.0.get());
    }

    #[test]
    #[cfg(panic = "unwind")]
    fn test_build_maximal_slots() {
        let builder = PyModuleSlotsBuilder::new()
            .with_mod_exec(module_exec)
            .with_name(c"test_module")
            .with_doc(c"some doc")
            .with_gil_used(false)
            .with_abi_info();

        assert!(builder.values[builder.len] == unsafe { std::mem::zeroed() });
        assert!(builder.values[builder.len - 1] != unsafe { std::mem::zeroed() });
        assert!(builder.len == MAX_SLOTS);

        let result = std::panic::catch_unwind(|| builder.with_mod_exec(module_exec).build());

        assert!(result.is_err());
    }

    #[test]
    #[should_panic]
    fn test_module_slots_builder_overflow() {
        let mut builder = PyModuleSlotsBuilder::new();
        for _ in 0..MAX_SLOTS + 1 {
            builder = builder.with_mod_exec(module_exec);
        }
    }
}