Update pyo3 to v0.22

Author: bschoenmaeckers

Cargo.toml

@@ -120,11 +120,11 @@ futures = "0.3"
 inventory = { version = "0.3", optional = true }
 once_cell = "1.14"
 pin-project-lite = "0.2"
-pyo3 = "0.21"
+pyo3 = "0.22"
 pyo3-async-runtimes-macros = { path = "pyo3-asyncio-macros", version = "=0.21.0", optional = true }
 
 [dev-dependencies]
-pyo3 = { version = "0.21", features = ["macros"] }
+pyo3 = { version = "0.22", features = ["macros"] }
 
 [dependencies.async-std]
 version = "1.12"

src/async_std.rs

@@ -13,11 +13,10 @@
 //! features = ["unstable-streams"]
 //! ```
 
-use std::{any::Any, cell::RefCell, future::Future, panic::AssertUnwindSafe, pin::Pin};
-
 use async_std::task;
 use futures::FutureExt;
 use pyo3::prelude::*;
+use std::{any::Any, cell::RefCell, future::Future, panic::AssertUnwindSafe, pin::Pin};
 
 use crate::{
     generic::{self, ContextExt, JoinError, LocalContextExt, Runtime, SpawnLocalExt},
@@ -90,10 +89,13 @@ impl ContextExt for AsyncStdRuntime {
     }
 
     fn get_task_locals() -> Option<TaskLocals> {
-        match TASK_LOCALS.try_with(|c| c.borrow().clone()) {
-            Ok(locals) => locals,
-            Err(_) => None,
-        }
+        TASK_LOCALS
+            .try_with(|c| {
+                c.borrow()
+                    .as_ref()
+                    .map(|locals| Python::with_gil(|py| locals.clone_ref(py)))
+            })
+            .unwrap_or_default()
     }
 }
 

src/generic.rs

@@ -15,24 +15,24 @@
 
 use std::{
     future::Future,
-    marker::PhantomData,
     pin::Pin,
     sync::{Arc, Mutex},
     task::{Context, Poll},
 };
 
-use futures::{
-    channel::{mpsc, oneshot},
-    SinkExt,
-};
-use once_cell::sync::OnceCell;
-use pin_project_lite::pin_project;
-use pyo3::prelude::*;
-
 use crate::{
     asyncio, call_soon_threadsafe, close, create_future, dump_err, err::RustPanic,
     get_running_loop, into_future_with_locals, TaskLocals,
 };
+use futures::channel::oneshot;
+#[cfg(feature = "unstable-streams")]
+use futures::{channel::mpsc, SinkExt};
+#[cfg(feature = "unstable-streams")]
+use once_cell::sync::OnceCell;
+use pin_project_lite::pin_project;
+use pyo3::prelude::*;
+#[cfg(feature = "unstable-streams")]
+use std::marker::PhantomData;
 
 /// Generic utilities for a JoinError
 pub trait JoinError {
@@ -585,7 +585,7 @@ where
 {
     let (cancel_tx, cancel_rx) = oneshot::channel();
 
-    let py_fut = create_future(locals.event_loop.clone().into_bound(py))?;
+    let py_fut = create_future(locals.event_loop.bind(py).clone())?;
     py_fut.call_method1(
         "add_done_callback",
         (PyDoneCallback {
@@ -594,14 +594,14 @@ where
     )?;
 
     let future_tx1 = PyObject::from(py_fut.clone());
-    let future_tx2 = future_tx1.clone();
+    let future_tx2 = future_tx1.clone_ref(py);
 
     R::spawn(async move {
-        let locals2 = locals.clone();
+        let locals2 = Python::with_gil(|py| locals.clone_ref(py));
 
         if let Err(e) = R::spawn(async move {
             let result = R::scope(
-                locals2.clone(),
+                Python::with_gil(|py| locals2.clone_ref(py)),
                 Cancellable::new_with_cancel_rx(fut, cancel_rx),
             )
             .await;
@@ -990,7 +990,7 @@ where
 {
     let (cancel_tx, cancel_rx) = oneshot::channel();
 
-    let py_fut = create_future(locals.event_loop.clone().into_bound(py))?;
+    let py_fut = create_future(locals.event_loop.clone_ref(py).into_bound(py))?;
     py_fut.call_method1(
         "add_done_callback",
         (PyDoneCallback {
@@ -999,14 +999,14 @@ where
     )?;
 
     let future_tx1 = PyObject::from(py_fut.clone());
-    let future_tx2 = future_tx1.clone();
+    let future_tx2 = future_tx1.clone_ref(py);
 
     R::spawn_local(async move {
-        let locals2 = locals.clone();
+        let locals2 = Python::with_gil(|py| locals.clone_ref(py));
 
         if let Err(e) = R::spawn_local(async move {
             let result = R::scope_local(
-                locals2.clone(),
+                Python::with_gil(|py| locals2.clone_ref(py)),
                 Cancellable::new_with_cancel_rx(fut, cancel_rx),
             )
             .await;
@@ -1446,27 +1446,12 @@ where
     into_stream_with_locals_v1::<R>(get_current_locals::<R>(gen.py())?, gen)
 }
 
-#[allow(dead_code)]
-fn py_true() -> PyObject {
-    static TRUE: OnceCell<PyObject> = OnceCell::new();
-    TRUE.get_or_init(|| Python::with_gil(|py| true.into_py(py)))
-        .clone()
-}
-
-#[allow(dead_code)]
-fn py_false() -> PyObject {
-    static FALSE: OnceCell<PyObject> = OnceCell::new();
-    FALSE
-        .get_or_init(|| Python::with_gil(|py| false.into_py(py)))
-        .clone()
-}
-
 trait Sender: Send + 'static {
-    fn send(&mut self, locals: TaskLocals, item: PyObject) -> PyResult<PyObject>;
+    fn send(&mut self, py: Python, locals: TaskLocals, item: PyObject) -> PyResult<PyObject>;
     fn close(&mut self) -> PyResult<()>;
 }
 
-#[allow(dead_code)]
+#[cfg(feature = "unstable-streams")]
 struct GenericSender<R>
 where
     R: Runtime,
@@ -1475,13 +1460,14 @@ where
     tx: mpsc::Sender<PyObject>,
 }
 
+#[cfg(feature = "unstable-streams")]
 impl<R> Sender for GenericSender<R>
 where
     R: Runtime + ContextExt,
 {
-    fn send(&mut self, locals: TaskLocals, item: PyObject) -> PyResult<PyObject> {
-        match self.tx.try_send(item.clone()) {
-            Ok(_) => Ok(py_true()),
+    fn send(&mut self, py: Python, locals: TaskLocals, item: PyObject) -> PyResult<PyObject> {
+        match self.tx.try_send(item.clone_ref(py)) {
+            Ok(_) => Ok(true.into_py(py)),
             Err(e) => {
                 if e.is_full() {
                     let mut tx = self.tx.clone();
@@ -1490,19 +1476,19 @@ where
                             future_into_py_with_locals::<R, _, PyObject>(py, locals, async move {
                                 if tx.flush().await.is_err() {
                                     // receiving side disconnected
-                                    return Ok(py_false());
+                                    return Python::with_gil(|py| Ok(false.into_py(py)));
                                 }
                                 if tx.send(item).await.is_err() {
                                     // receiving side disconnected
-                                    return Ok(py_false());
+                                    return Python::with_gil(|py| Ok(false.into_py(py)));
                                 }
-                                Ok(py_true())
+                                Python::with_gil(|py| Ok(true.into_py(py)))
                             })?
                             .into(),
                         )
                     })
                 } else {
-                    Ok(py_false())
+                    Ok(false.into_py(py))
                 }
             }
         }
@@ -1521,7 +1507,7 @@ struct SenderGlue {
 #[pymethods]
 impl SenderGlue {
     pub fn send(&mut self, item: PyObject) -> PyResult<PyObject> {
-        self.tx.send(self.locals.clone(), item)
+        Python::with_gil(|py| self.tx.send(py, self.locals.clone_ref(py), item))
     }
     pub fn close(&mut self) -> PyResult<()> {
         self.tx.close()

src/lib.rs

@@ -473,7 +473,7 @@ fn copy_context(py: Python) -> PyResult<Bound<PyAny>> {
 }
 
 /// Task-local data to store for Python conversions.
-#[derive(Debug, Clone)]
+#[derive(Debug)]
 pub struct TaskLocals {
     /// Track the event loop of the Python task
     event_loop: PyObject,
@@ -510,12 +510,21 @@ impl TaskLocals {
 
     /// Get a reference to the event loop
     pub fn event_loop<'p>(&self, py: Python<'p>) -> Bound<'p, PyAny> {
-        self.event_loop.clone().into_bound(py)
+        self.event_loop.clone_ref(py).into_bound(py)
     }
 
     /// Get a reference to the python context
     pub fn context<'p>(&self, py: Python<'p>) -> Bound<'p, PyAny> {
-        self.context.clone().into_bound(py)
+        self.context.clone_ref(py).into_bound(py)
+    }
+
+    /// Create a clone of the TaskLocals by incrementing the reference counters of the event loop and
+    /// contextvars.
+    pub fn clone_ref(&self, py: Python<'_>) -> Self {
+        Self {
+            event_loop: self.event_loop.clone_ref(py),
+            context: self.context.clone_ref(py),
+        }
     }
 }
 

src/tokio.rs

@@ -109,10 +109,12 @@ impl ContextExt for TokioRuntime {
     }
 
     fn get_task_locals() -> Option<TaskLocals> {
-        match TASK_LOCALS.try_with(|c| c.get().map(|locals| locals.clone())) {
-            Ok(locals) => locals,
-            Err(_) => None,
-        }
+        TASK_LOCALS
+            .try_with(|c| {
+                c.get()
+                    .map(|locals| Python::with_gil(|py| locals.clone_ref(py)))
+            })
+            .unwrap_or_default()
     }
 }