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}