From d7e147509bbdf3b21353b71f9f60cdbad54acc7f Mon Sep 17 00:00:00 2001
From: Adam Reichold <adam.reichold@t-online.de>
Date: Wed, 21 Jun 2023 08:00:26 +0200
Subject: [PATCH] Add Python::with_pool as a safer alternative to
 Python::new_pool.

---
 newsfragments/3263.added.md |   1 +
 src/marker.rs               | 131 ++++++++++++++++++++++++++++++++++++
 2 files changed, 132 insertions(+)
 create mode 100644 newsfragments/3263.added.md

diff --git a/newsfragments/3263.added.md b/newsfragments/3263.added.md
new file mode 100644
index 00000000..888395d6
--- /dev/null
+++ b/newsfragments/3263.added.md
@@ -0,0 +1 @@
+Add the `Python::with_pool` which is a safer but more limited alternative to `Python::new_pool`.
diff --git a/src/marker.rs b/src/marker.rs
index 129ed367..da592701 100644
--- a/src/marker.rs
+++ b/src/marker.rs
@@ -940,6 +940,137 @@ impl<'py> Python<'py> {
     }
 }
 
+impl Python<'_> {
+    /// Creates a scope using a new pool for managing PyO3's owned references.
+    ///
+    /// This is a safe alterantive to [`new_pool`][Self::new_pool] as
+    /// it limits the closure to using the new GIL token at the cost of
+    /// being unable to capture existing GIL-bound references.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// # use pyo3::prelude::*;
+    /// Python::with_gil(|py| {
+    ///     // Some long-running process like a webserver, which never releases the GIL.
+    ///     loop {
+    ///         // Create a new scope, so that PyO3 can clear memory at the end of the loop.
+    ///         py.with_pool(|py| {
+    ///             // do stuff...
+    ///         });
+    /// #       break;  // Exit the loop so that doctest terminates!
+    ///     }
+    /// });
+    /// ```
+    ///
+    /// The `Ungil` bound on the closure does prevent hanging on to existing GIL-bound references
+    ///
+    /// ```compile_fail
+    /// # use pyo3::prelude::*;
+    /// # use pyo3::types::PyString;
+    ///
+    /// Python::with_gil(|py| {
+    ///     let old_str = PyString::new(py, "a message from the past");
+    ///
+    ///     py.with_pool(|_py| {
+    ///         print!("{:?}", old_str);
+    ///     });
+    /// });
+    /// ```
+    ///
+    /// or continuing to use the old GIL token
+    ///
+    /// ```compile_fail
+    /// # use pyo3::prelude::*;
+    ///
+    /// Python::with_gil(|old_py| {
+    ///     old_py.with_pool(|_new_py| {
+    ///         let _none = old_py.None();
+    ///     });
+    /// });
+    /// ```
+    #[cfg(feature = "nightly")]
+    #[inline]
+    pub fn with_pool<F, R>(&self, f: F) -> R
+    where
+        F: for<'py> FnOnce(Python<'py>) -> R + Ungil,
+    {
+        // SAFETY: The closure is `Ungil`,
+        // i.e. it does not capture any GIL-bound references
+        // and accesses only the newly created GIL token.
+        let pool = unsafe { GILPool::new() };
+
+        f(pool.python())
+    }
+
+    /// Creates a scope using a new pool for managing PyO3's owned references.
+    ///
+    /// This is a safer alterantive to [`new_pool`][Self::new_pool] as
+    /// it limits the closure to using the new GIL token at the cost of
+    /// being unable to capture existing GIL-bound references.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// # use pyo3::prelude::*;
+    /// Python::with_gil(|py| {
+    ///     // Some long-running process like a webserver, which never releases the GIL.
+    ///     loop {
+    ///         // Create a new scope, so that PyO3 can clear memory at the end of the loop.
+    ///         unsafe {
+    ///             py.with_pool(|py| {
+    ///                 // do stuff...
+    ///             });
+    ///         }
+    /// #       break;  // Exit the loop so that doctest terminates!
+    ///     }
+    /// });
+    /// ```
+    ///
+    /// The `Ungil` bound on the closure does prevent hanging on to existing GIL-bound references
+    ///
+    /// ```compile_fail
+    /// # use pyo3::prelude::*;
+    /// # use pyo3::types::PyString;
+    ///
+    /// Python::with_gil(|py| {
+    ///     let old_str = PyString::new(py, "a message from the past");
+    ///
+    ///     py.with_pool(|_py| {
+    ///         print!("{:?}", old_str);
+    ///     });
+    /// });
+    /// ```
+    ///
+    /// or continuing to use the old GIL token
+    ///
+    /// ```compile_fail
+    /// # use pyo3::prelude::*;
+    ///
+    /// Python::with_gil(|old_py| {
+    ///     old_py.with_pool(|_new_py| {
+    ///         let _none = old_py.None();
+    ///     });
+    /// });
+    /// ```
+    ///
+    /// # Safety
+    ///
+    /// However, due to the `SendWrapper` loophole describe in the documentation of the [`Ungil`] trait,
+    /// this is still an unsafe API. Usage that does not involve runtime checking of thread affinity
+    /// should be practically safe.
+    #[cfg(not(feature = "nightly"))]
+    #[inline]
+    pub unsafe fn with_pool<F, R>(&self, f: F) -> R
+    where
+        F: for<'py> FnOnce(Python<'py>) -> R + Ungil,
+    {
+        let pool = GILPool::new();
+
+        f(pool.python())
+    }
+}
+
 impl<'unbound> Python<'unbound> {
     /// Unsafely creates a Python token with an unbounded lifetime.
     ///