polars_python/
gil_once_cell.rs

1use std::cell::UnsafeCell;
2
3use pyo3::{PyResult, Python};
4
5// Adapted from PYO3 with the only change that
6// we allow mutable access with when the GIL is held
7
8pub struct GILOnceCell<T>(UnsafeCell<Option<T>>);
9
10// T: Send is needed for Sync because the thread which drops the GILOnceCell can be different
11// to the thread which fills it.
12unsafe impl<T: Send + Sync> Sync for GILOnceCell<T> {}
13unsafe impl<T: Send> Send for GILOnceCell<T> {}
14
15impl<T> GILOnceCell<T> {
16    /// Create a `GILOnceCell` which does not yet contain a value.
17    #[allow(clippy::new_without_default)]
18    pub const fn new() -> Self {
19        Self(UnsafeCell::new(None))
20    }
21
22    /// as long as we have the GIL we can mutate
23    /// this creates a context that checks that.
24    pub fn with_gil<F, O>(&self, _py: Python<'_>, mut op: F) -> PyResult<O>
25    where
26        F: FnMut(&mut T) -> PyResult<O>,
27    {
28        // Safe because GIL is held, so no other thread can be writing to this cell concurrently.
29        let inner = unsafe { &mut *self.0.get() }
30            .as_mut()
31            .expect("not yet initialized");
32
33        op(inner)
34    }
35
36    /// Set the value in the cell.
37    ///
38    /// If the cell has already been written, `Err(value)` will be returned containing the new
39    /// value which was not written.
40    pub fn set(&self, _py: Python<'_>, value: T) -> Result<(), T> {
41        // Safe because GIL is held, so no other thread can be writing to this cell concurrently.
42        let inner = unsafe { &mut *self.0.get() };
43        if inner.is_some() {
44            return Err(value);
45        }
46
47        *inner = Some(value);
48        Ok(())
49    }
50}