From 632ff756d8b851bc05677432848bbf7cd77b66f5 Mon Sep 17 00:00:00 2001
From: Nathan Goldbaum <nathan.goldbaum@gmail.com>
Date: Tue, 3 Sep 2024 14:06:49 -0600
Subject: [PATCH] add pyo3::types::mutex for PyMutex wrappers

---
 src/types/mod.rs   |  3 ++
 src/types/mutex.rs | 86 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 89 insertions(+)
 create mode 100644 src/types/mutex.rs

diff --git a/src/types/mod.rs b/src/types/mod.rs
index d11af8598fe..5d18b99b785 100644
--- a/src/types/mod.rs
+++ b/src/types/mod.rs
@@ -30,6 +30,7 @@ pub use self::list::{PyList, PyListMethods};
 pub use self::mapping::{PyMapping, PyMappingMethods};
 pub use self::memoryview::PyMemoryView;
 pub use self::module::{PyModule, PyModuleMethods};
+pub use self::mutex::{PyMutex, PyMutexGuard};
 pub use self::none::PyNone;
 pub use self::notimplemented::PyNotImplemented;
 #[allow(deprecated)]
@@ -243,6 +244,8 @@ pub(crate) mod list;
 pub(crate) mod mapping;
 mod memoryview;
 pub(crate) mod module;
+#[cfg(Py_3_13)]
+mod mutex;
 mod none;
 mod notimplemented;
 mod num;
diff --git a/src/types/mutex.rs b/src/types/mutex.rs
new file mode 100644
index 00000000000..c0c139159ea
--- /dev/null
+++ b/src/types/mutex.rs
@@ -0,0 +1,86 @@
+use std::ops::Deref;
+
+/// Wrapper for [`PyMutex`](https://docs.python.org/3/c-api/init.html#c.PyMutex) exposing an RAII interface.
+#[derive(Debug)]
+pub struct PyMutex<T> {
+    _mutex: crate::ffi::PyMutex,
+    data: T,
+}
+
+/// RAII guard to handle releasing a PyMutex lock.
+#[derive(Debug)]
+pub struct PyMutexGuard<'a, T> {
+    _mutex: &'a mut crate::ffi::PyMutex,
+    data: &'a T,
+}
+
+impl<T> PyMutex<T> {
+    /// Acquire the mutex, blocking the current thread until it is able to do so.
+    pub fn lock(&mut self) -> PyMutexGuard<'_, T> {
+        unsafe { crate::ffi::PyMutex_Lock(&mut self._mutex) };
+        PyMutexGuard {
+            _mutex: &mut self._mutex,
+            data: &self.data,
+        }
+    }
+
+    /// Create a new mutex in an unlocked state ready for use.
+    pub fn new(value: T) -> Self {
+        Self {
+            _mutex: crate::ffi::PyMutex::new(),
+            data: value,
+        }
+    }
+}
+
+impl<'a, T> Drop for PyMutexGuard<'a, T> {
+    fn drop(&mut self) {
+        unsafe { crate::ffi::PyMutex_Unlock(self._mutex) };
+    }
+}
+
+impl<'a, T> Deref for PyMutexGuard<'a, T> {
+    type Target = T;
+
+    fn deref(&self) -> &T {
+        self.data
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::types::{PyAnyMethods, PyDict, PyDictMethods, PyList, PyNone};
+    use crate::Python;
+
+    #[test]
+    fn test_pymutex() {
+        Python::with_gil(|py| {
+            let d = PyDict::new(py);
+            let mut mutex = PyMutex::new(&d);
+
+            let list = Python::with_gil(|py| PyList::new(py, vec!["foo", "bar"]).unbind());
+            let dict_guard = mutex.lock();
+
+            py.allow_threads(|| {
+                std::thread::spawn(move || {
+                    drop(list);
+                })
+                .join()
+                .unwrap();
+            });
+
+            dict_guard
+                .set_item(PyNone::get(py), PyNone::get(py))
+                .unwrap();
+            drop(dict_guard);
+
+            assert!(d
+                .get_item(PyNone::get(py))
+                .unwrap()
+                .unwrap()
+                .eq(PyNone::get(py))
+                .unwrap());
+        });
+    }
+}