pyo3/impl_/
pymodule.rs

Help
1//! Implementation details of `#[pymodule]` which need to be accessible from proc-macro generated code.
2
3use std::{cell::UnsafeCell, ffi::CStr, marker::PhantomData};
4
5#[cfg(all(
6    not(any(PyPy, GraalPy)),
7    Py_3_9,
8    not(all(windows, Py_LIMITED_API, not(Py_3_10))),
9    not(target_has_atomic = "64"),
10))]
11use portable_atomic::AtomicI64;
12#[cfg(all(
13    not(any(PyPy, GraalPy)),
14    Py_3_9,
15    not(all(windows, Py_LIMITED_API, not(Py_3_10))),
16    target_has_atomic = "64",
17))]
18use std::sync::atomic::AtomicI64;
19use std::sync::atomic::{AtomicBool, Ordering};
20
21#[cfg(not(any(PyPy, GraalPy)))]
22use crate::exceptions::PyImportError;
23#[cfg(all(not(Py_LIMITED_API), Py_GIL_DISABLED))]
24use crate::PyErr;
25use crate::{
26    ffi,
27    impl_::pymethods::PyMethodDef,
28    sync::GILOnceCell,
29    types::{PyCFunction, PyModule, PyModuleMethods},
30    Bound, Py, PyClass, PyResult, PyTypeInfo, Python,
31};
32
33/// `Sync` wrapper of `ffi::PyModuleDef`.
34pub struct ModuleDef {
35    // wrapped in UnsafeCell so that Rust compiler treats this as interior mutability
36    ffi_def: UnsafeCell<ffi::PyModuleDef>,
37    initializer: ModuleInitializer,
38    /// Interpreter ID where module was initialized (not applicable on PyPy).
39    #[cfg(all(
40        not(any(PyPy, GraalPy)),
41        Py_3_9,
42        not(all(windows, Py_LIMITED_API, not(Py_3_10)))
43    ))]
44    interpreter: AtomicI64,
45    /// Initialized module object, cached to avoid reinitialization.
46    module: GILOnceCell<Py<PyModule>>,
47    /// Whether or not the module supports running without the GIL
48    gil_used: AtomicBool,
49}
50
51/// Wrapper to enable initializer to be used in const fns.
52pub struct ModuleInitializer(pub for<'py> fn(&Bound<'py, PyModule>) -> PyResult<()>);
53
54unsafe impl Sync for ModuleDef {}
55
56impl ModuleDef {
57    /// Make new module definition with given module name.
58    pub const unsafe fn new(
59        name: &'static CStr,
60        doc: &'static CStr,
61        initializer: ModuleInitializer,
62    ) -> Self {
63        #[allow(clippy::declare_interior_mutable_const)]
64        const INIT: ffi::PyModuleDef = ffi::PyModuleDef {
65            m_base: ffi::PyModuleDef_HEAD_INIT,
66            m_name: std::ptr::null(),
67            m_doc: std::ptr::null(),
68            m_size: 0,
69            m_methods: std::ptr::null_mut(),
70            m_slots: std::ptr::null_mut(),
71            m_traverse: None,
72            m_clear: None,
73            m_free: None,
74        };
75
76        let ffi_def = UnsafeCell::new(ffi::PyModuleDef {
77            m_name: name.as_ptr(),
78            m_doc: doc.as_ptr(),
79            ..INIT
80        });
81
82        ModuleDef {
83            ffi_def,
84            initializer,
85            // -1 is never expected to be a valid interpreter ID
86            #[cfg(all(
87                not(any(PyPy, GraalPy)),
88                Py_3_9,
89                not(all(windows, Py_LIMITED_API, not(Py_3_10)))
90            ))]
91            interpreter: AtomicI64::new(-1),
92            module: GILOnceCell::new(),
93            gil_used: AtomicBool::new(true),
94        }
95    }
96    /// Builds a module using user given initializer. Used for [`#[pymodule]`][crate::pymodule].
97    #[cfg_attr(any(Py_LIMITED_API, not(Py_GIL_DISABLED)), allow(unused_variables))]
98    pub fn make_module(&'static self, py: Python<'_>, gil_used: bool) -> PyResult<Py<PyModule>> {
99        // Check the interpreter ID has not changed, since we currently have no way to guarantee
100        // that static data is not reused across interpreters.
101        //
102        // PyPy does not have subinterpreters, so no need to check interpreter ID.
103        #[cfg(not(any(PyPy, GraalPy)))]
104        {
105            // PyInterpreterState_Get is only available on 3.9 and later, but is missing
106            // from python3.dll for Windows stable API on 3.9
107            #[cfg(all(Py_3_9, not(all(windows, Py_LIMITED_API, not(Py_3_10)))))]
108            {
109                let current_interpreter =
110                    unsafe { ffi::PyInterpreterState_GetID(ffi::PyInterpreterState_Get()) };
111                crate::err::error_on_minusone(py, current_interpreter)?;
112                if let Err(initialized_interpreter) = self.interpreter.compare_exchange(
113                    -1,
114                    current_interpreter,
115                    Ordering::SeqCst,
116                    Ordering::SeqCst,
117                ) {
118                    if initialized_interpreter != current_interpreter {
119                        return Err(PyImportError::new_err(
120                            "PyO3 modules do not yet support subinterpreters, see https://github.com/PyO3/pyo3/issues/576",
121                        ));
122                    }
123                }
124            }
125            #[cfg(not(all(Py_3_9, not(all(windows, Py_LIMITED_API, not(Py_3_10))))))]
126            {
127                // CPython before 3.9 does not have APIs to check the interpreter ID, so best that can be
128                // done to guard against subinterpreters is fail if the module is initialized twice
129                if self.module.get(py).is_some() {
130                    return Err(PyImportError::new_err(
131                        "PyO3 modules compiled for CPython 3.8 or older may only be initialized once per interpreter process"
132                    ));
133                }
134            }
135        }
136        self.module
137            .get_or_try_init(py, || {
138                let module = unsafe {
139                    Py::<PyModule>::from_owned_ptr_or_err(
140                        py,
141                        ffi::PyModule_Create(self.ffi_def.get()),
142                    )?
143                };
144                #[cfg(all(not(Py_LIMITED_API), Py_GIL_DISABLED))]
145                {
146                    let gil_used_ptr = {
147                        if gil_used {
148                            ffi::Py_MOD_GIL_USED
149                        } else {
150                            ffi::Py_MOD_GIL_NOT_USED
151                        }
152                    };
153                    if unsafe { ffi::PyUnstable_Module_SetGIL(module.as_ptr(), gil_used_ptr) } < 0 {
154                        return Err(PyErr::fetch(py));
155                    }
156                }
157                self.initializer.0(module.bind(py))?;
158                Ok(module)
159            })
160            .map(|py_module| py_module.clone_ref(py))
161    }
162}
163
164/// Trait to add an element (class, function...) to a module.
165///
166/// Currently only implemented for classes.
167pub trait PyAddToModule: crate::sealed::Sealed {
168    fn add_to_module(&'static self, module: &Bound<'_, PyModule>) -> PyResult<()>;
169}
170
171/// For adding native types (non-pyclass) to a module.
172pub struct AddTypeToModule<T>(PhantomData<T>);
173
174impl<T> AddTypeToModule<T> {
175    #[allow(clippy::new_without_default)]
176    pub const fn new() -> Self {
177        AddTypeToModule(PhantomData)
178    }
179}
180
181impl<T: PyTypeInfo> PyAddToModule for AddTypeToModule<T> {
182    fn add_to_module(&'static self, module: &Bound<'_, PyModule>) -> PyResult<()> {
183        module.add(T::NAME, T::type_object(module.py()))
184    }
185}
186
187/// For adding a class to a module.
188pub struct AddClassToModule<T>(PhantomData<T>);
189
190impl<T> AddClassToModule<T> {
191    #[allow(clippy::new_without_default)]
192    pub const fn new() -> Self {
193        AddClassToModule(PhantomData)
194    }
195}
196
197impl<T: PyClass> PyAddToModule for AddClassToModule<T> {
198    fn add_to_module(&'static self, module: &Bound<'_, PyModule>) -> PyResult<()> {
199        module.add_class::<T>()
200    }
201}
202
203/// For adding a function to a module.
204impl PyAddToModule for PyMethodDef {
205    fn add_to_module(&'static self, module: &Bound<'_, PyModule>) -> PyResult<()> {
206        module.add_function(PyCFunction::internal_new(module.py(), self, Some(module))?)
207    }
208}
209
210/// For adding a module to a module.
211impl PyAddToModule for ModuleDef {
212    fn add_to_module(&'static self, module: &Bound<'_, PyModule>) -> PyResult<()> {
213        module.add_submodule(
214            self.make_module(module.py(), self.gil_used.load(Ordering::Relaxed))?
215                .bind(module.py()),
216        )
217    }
218}
219
220#[cfg(test)]
221mod tests {
222    use std::{
223        borrow::Cow,
224        ffi::CStr,
225        sync::atomic::{AtomicBool, Ordering},
226    };
227
228    use crate::{
229        ffi,
230        types::{any::PyAnyMethods, module::PyModuleMethods, PyModule},
231        Bound, PyResult, Python,
232    };
233
234    use super::{ModuleDef, ModuleInitializer};
235
236    #[test]
237    fn module_init() {
238        static MODULE_DEF: ModuleDef = unsafe {
239            ModuleDef::new(
240                ffi::c_str!("test_module"),
241                ffi::c_str!("some doc"),
242                ModuleInitializer(|m| {
243                    m.add("SOME_CONSTANT", 42)?;
244                    Ok(())
245                }),
246            )
247        };
248        Python::with_gil(|py| {
249            let module = MODULE_DEF.make_module(py, false).unwrap().into_bound(py);
250            assert_eq!(
251                module
252                    .getattr("__name__")
253                    .unwrap()
254                    .extract::<Cow<'_, str>>()
255                    .unwrap(),
256                "test_module",
257            );
258            assert_eq!(
259                module
260                    .getattr("__doc__")
261                    .unwrap()
262                    .extract::<Cow<'_, str>>()
263                    .unwrap(),
264                "some doc",
265            );
266            assert_eq!(
267                module
268                    .getattr("SOME_CONSTANT")
269                    .unwrap()
270                    .extract::<u8>()
271                    .unwrap(),
272                42,
273            );
274        })
275    }
276
277    #[test]
278    fn module_def_new() {
279        // To get coverage for ModuleDef::new() need to create a non-static ModuleDef, however init
280        // etc require static ModuleDef, so this test needs to be separated out.
281        static NAME: &CStr = ffi::c_str!("test_module");
282        static DOC: &CStr = ffi::c_str!("some doc");
283
284        static INIT_CALLED: AtomicBool = AtomicBool::new(false);
285
286        #[allow(clippy::unnecessary_wraps)]
287        fn init(_: &Bound<'_, PyModule>) -> PyResult<()> {
288            INIT_CALLED.store(true, Ordering::SeqCst);
289            Ok(())
290        }
291
292        unsafe {
293            let module_def: ModuleDef = ModuleDef::new(NAME, DOC, ModuleInitializer(init));
294            assert_eq!((*module_def.ffi_def.get()).m_name, NAME.as_ptr() as _);
295            assert_eq!((*module_def.ffi_def.get()).m_doc, DOC.as_ptr() as _);
296
297            Python::with_gil(|py| {
298                module_def.initializer.0(&py.import("builtins").unwrap()).unwrap();
299                assert!(INIT_CALLED.load(Ordering::SeqCst));
300            })
301        }
302    }
303}
pyo3/impl_/pymodule.rs

pyo3/impl_/
pymodule.rs
