pyo3/sync/
once_lock.rs

1use crate::{
2    internal::state::SuspendAttach, types::any::PyAnyMethods, Bound, Py, PyResult, PyTypeCheck,
3    Python,
4};
5
6/// An equivalent to [`std::sync::OnceLock`] for initializing objects while attached to
7/// the Python interpreter.
8///
9/// Unlike `OnceLock<T>`, this type will not deadlock with the interpreter.
10/// Before blocking calls the cell will detach from the runtime and then
11/// re-attach once the cell is unblocked.
12///
13/// # Re-entrant initialization
14///
15/// Like `OnceLock<T>`, it is an error to re-entrantly initialize a `PyOnceLock<T>`. The exact
16/// behavior in this case is not guaranteed, it may either deadlock or panic.
17///
18/// # Examples
19///
20/// The following example shows how to use `PyOnceLock` to share a reference to a Python list
21/// between threads:
22///
23/// ```
24/// use pyo3::sync::PyOnceLock;
25/// use pyo3::prelude::*;
26/// use pyo3::types::PyList;
27///
28/// static LIST_CELL: PyOnceLock<Py<PyList>> = PyOnceLock::new();
29///
30/// pub fn get_shared_list(py: Python<'_>) -> &Bound<'_, PyList> {
31///     LIST_CELL
32///         .get_or_init(py, || PyList::empty(py).unbind())
33///         .bind(py)
34/// }
35/// # Python::attach(|py| assert_eq!(get_shared_list(py).len(), 0));
36/// ```
37#[derive(Default)]
38pub struct PyOnceLock<T> {
39    inner: once_cell::sync::OnceCell<T>,
40}
41
42impl<T> PyOnceLock<T> {
43    /// Create a `PyOnceLock` which does not yet contain a value.
44    pub const fn new() -> Self {
45        Self {
46            inner: once_cell::sync::OnceCell::new(),
47        }
48    }
49
50    /// Get a reference to the contained value, or `None` if the cell has not yet been written.
51    #[inline]
52    pub fn get(&self, _py: Python<'_>) -> Option<&T> {
53        self.inner.get()
54    }
55
56    /// Get a reference to the contained value, initializing it if needed using the provided
57    /// closure.
58    ///
59    /// See the type-level documentation for detail on re-entrancy and concurrent initialization.
60    #[inline]
61    pub fn get_or_init<F>(&self, py: Python<'_>, f: F) -> &T
62    where
63        F: FnOnce() -> T,
64    {
65        self.inner
66            .get()
67            .unwrap_or_else(|| init_once_cell_py_attached(&self.inner, py, f))
68    }
69
70    /// Like `get_or_init`, but accepts a fallible initialization function. If it fails, the cell
71    /// is left uninitialized.
72    ///
73    /// See the type-level documentation for detail on re-entrancy and concurrent initialization.
74    pub fn get_or_try_init<F, E>(&self, py: Python<'_>, f: F) -> Result<&T, E>
75    where
76        F: FnOnce() -> Result<T, E>,
77    {
78        self.inner
79            .get()
80            .map_or_else(|| try_init_once_cell_py_attached(&self.inner, py, f), Ok)
81    }
82
83    /// Get the contents of the cell mutably. This is only possible if the reference to the cell is
84    /// unique.
85    pub fn get_mut(&mut self) -> Option<&mut T> {
86        self.inner.get_mut()
87    }
88
89    /// Set the value in the cell.
90    ///
91    /// If the cell has already been written, `Err(value)` will be returned containing the new
92    /// value which was not written.
93    pub fn set(&self, _py: Python<'_>, value: T) -> Result<(), T> {
94        self.inner.set(value)
95    }
96
97    /// Takes the value out of the cell, moving it back to an uninitialized state.
98    ///
99    /// Has no effect and returns None if the cell has not yet been written.
100    pub fn take(&mut self) -> Option<T> {
101        self.inner.take()
102    }
103
104    /// Consumes the cell, returning the wrapped value.
105    ///
106    /// Returns None if the cell has not yet been written.
107    pub fn into_inner(self) -> Option<T> {
108        self.inner.into_inner()
109    }
110}
111
112impl<T> PyOnceLock<Py<T>> {
113    /// Creates a new cell that contains a new Python reference to the same contained object.
114    ///
115    /// Returns an uninitialized cell if `self` has not yet been initialized.
116    pub fn clone_ref(&self, py: Python<'_>) -> Self {
117        let cloned = PyOnceLock::new();
118        if let Some(value) = self.get(py) {
119            let _ = cloned.set(py, value.clone_ref(py));
120        }
121        cloned
122    }
123}
124
125impl<T> PyOnceLock<Py<T>>
126where
127    T: PyTypeCheck,
128{
129    /// This is a shorthand method for `get_or_init` which imports the type from Python on init.
130    ///
131    /// # Example: Using `PyOnceLock` to store a class in a static variable.
132    ///
133    /// `PyOnceLock` can be used to avoid importing a class multiple times:
134    /// ```
135    /// # use pyo3::prelude::*;
136    /// # use pyo3::sync::PyOnceLock;
137    /// # use pyo3::types::{PyDict, PyType};
138    /// # use pyo3::intern;
139    /// #
140    /// #[pyfunction]
141    /// fn create_ordered_dict<'py>(py: Python<'py>, dict: Bound<'py, PyDict>) -> PyResult<Bound<'py, PyAny>> {
142    ///     // Even if this function is called multiple times,
143    ///     // the `OrderedDict` class will be imported only once.
144    ///     static ORDERED_DICT: PyOnceLock<Py<PyType>> = PyOnceLock::new();
145    ///     ORDERED_DICT
146    ///         .import(py, "collections", "OrderedDict")?
147    ///         .call1((dict,))
148    /// }
149    ///
150    /// # Python::attach(|py| {
151    /// #     let dict = PyDict::new(py);
152    /// #     dict.set_item(intern!(py, "foo"), 42).unwrap();
153    /// #     let fun = wrap_pyfunction!(create_ordered_dict, py).unwrap();
154    /// #     let ordered_dict = fun.call1((&dict,)).unwrap();
155    /// #     assert!(dict.eq(ordered_dict).unwrap());
156    /// # });
157    /// ```
158    pub fn import<'py>(
159        &self,
160        py: Python<'py>,
161        module_name: &str,
162        attr_name: &str,
163    ) -> PyResult<&Bound<'py, T>> {
164        self.get_or_try_init(py, || {
165            let type_object = py
166                .import(module_name)?
167                .getattr(attr_name)?
168                .downcast_into()?;
169            Ok(type_object.unbind())
170        })
171        .map(|ty| ty.bind(py))
172    }
173}
174
175#[cold]
176fn init_once_cell_py_attached<'a, F, T>(
177    cell: &'a once_cell::sync::OnceCell<T>,
178    _py: Python<'_>,
179    f: F,
180) -> &'a T
181where
182    F: FnOnce() -> T,
183{
184    // SAFETY: detach from the runtime right before a possibly blocking call
185    // then reattach when the blocking call completes and before calling
186    // into the C API.
187    let ts_guard = unsafe { SuspendAttach::new() };
188
189    // By having detached here, we guarantee that `.get_or_init` cannot deadlock with
190    // the Python interpreter
191    cell.get_or_init(move || {
192        drop(ts_guard);
193        f()
194    })
195}
196
197#[cold]
198fn try_init_once_cell_py_attached<'a, F, T, E>(
199    cell: &'a once_cell::sync::OnceCell<T>,
200    _py: Python<'_>,
201    f: F,
202) -> Result<&'a T, E>
203where
204    F: FnOnce() -> Result<T, E>,
205{
206    // SAFETY: detach from the runtime right before a possibly blocking call
207    // then reattach when the blocking call completes and before calling
208    // into the C API.
209    let ts_guard = unsafe { SuspendAttach::new() };
210
211    // By having detached here, we guarantee that `.get_or_init` cannot deadlock with
212    // the Python interpreter
213    cell.get_or_try_init(move || {
214        drop(ts_guard);
215        f()
216    })
217}
218
219#[cfg(test)]
220mod tests {
221    use super::*;
222
223    #[test]
224    fn test_once_cell() {
225        Python::attach(|py| {
226            let mut cell = PyOnceLock::new();
227
228            assert!(cell.get(py).is_none());
229
230            assert_eq!(cell.get_or_try_init(py, || Err(5)), Err(5));
231            assert!(cell.get(py).is_none());
232
233            assert_eq!(cell.get_or_try_init(py, || Ok::<_, ()>(2)), Ok(&2));
234            assert_eq!(cell.get(py), Some(&2));
235
236            assert_eq!(cell.get_or_try_init(py, || Err(5)), Ok(&2));
237
238            assert_eq!(cell.take(), Some(2));
239            assert_eq!(cell.into_inner(), None);
240
241            let cell_py = PyOnceLock::new();
242            assert!(cell_py.clone_ref(py).get(py).is_none());
243            cell_py.get_or_init(py, || py.None());
244            assert!(cell_py.clone_ref(py).get(py).unwrap().is_none(py));
245        })
246    }
247
248    #[test]
249    fn test_once_cell_drop() {
250        #[derive(Debug)]
251        struct RecordDrop<'a>(&'a mut bool);
252
253        impl Drop for RecordDrop<'_> {
254            fn drop(&mut self) {
255                *self.0 = true;
256            }
257        }
258
259        Python::attach(|py| {
260            let mut dropped = false;
261            let cell = PyOnceLock::new();
262            cell.set(py, RecordDrop(&mut dropped)).unwrap();
263            let drop_container = cell.get(py).unwrap();
264
265            assert!(!*drop_container.0);
266            drop(cell);
267            assert!(dropped);
268        });
269    }
270}