Introduce #[pyclass(unsendable)]

guide/src/class.md

@@ -974,6 +974,10 @@ impl pyo3::class::proto_methods::HasProtoRegistry for MyClass {
         &REGISTRY
     }
 }
+
+unsafe impl pyo3::pyclass::PyClassSend for MyClass {
+    type ThreadChecker = pyo3::pyclass::ThreadCheckerStub<MyClass>;
+}
 # let gil = Python::acquire_gil();
 # let py = gil.python();
 # let cls = py.get_type::<MyClass>();

pyo3-derive-backend/src/pyclass.rs

@@ -19,6 +19,7 @@ pub struct PyClassArgs {
     pub flags: Vec<syn::Expr>,
     pub base: syn::TypePath,
     pub has_extends: bool,
+    pub has_unsendable: bool,
     pub module: Option<syn::LitStr>,
 }
 
@@ -45,6 +46,7 @@ impl Default for PyClassArgs {
             flags: vec![parse_quote! { 0 }],
             base: parse_quote! { pyo3::PyAny },
             has_extends: false,
+            has_unsendable: false,
         }
     }
 }
@@ -60,7 +62,7 @@ impl PyClassArgs {
         }
     }
 
-    /// Match a single flag
+    /// Match a key/value flag
     fn add_assign(&mut self, assign: &syn::ExprAssign) -> syn::Result<()> {
         let syn::ExprAssign { left, right, .. } = assign;
         let key = match &**left {
@@ -120,31 +122,27 @@ impl PyClassArgs {
         Ok(())
     }
 
-    /// Match a key/value flag
+    /// Match a single flag
     fn add_path(&mut self, exp: &syn::ExprPath) -> syn::Result<()> {
         let flag = exp.path.segments.first().unwrap().ident.to_string();
-        let path = match flag.as_str() {
-            "gc" => {
-                parse_quote! {pyo3::type_flags::GC}
-            }
-            "weakref" => {
-                parse_quote! {pyo3::type_flags::WEAKREF}
-            }
-            "subclass" => {
-                parse_quote! {pyo3::type_flags::BASETYPE}
-            }
-            "dict" => {
-                parse_quote! {pyo3::type_flags::DICT}
+        let mut push_flag = |flag| {
+            self.flags.push(syn::Expr::Path(flag));
+        };
+        match flag.as_str() {
+            "gc" => push_flag(parse_quote! {pyo3::type_flags::GC}),
+            "weakref" => push_flag(parse_quote! {pyo3::type_flags::WEAKREF}),
+            "subclass" => push_flag(parse_quote! {pyo3::type_flags::BASETYPE}),
+            "dict" => push_flag(parse_quote! {pyo3::type_flags::DICT}),
+            "unsendable" => {
+                self.has_unsendable = true;
             }
             _ => {
                 return Err(syn::Error::new_spanned(
                     &exp.path,
-                    "Expected one of gc/weakref/subclass/dict",
+                    "Expected one of gc/weakref/subclass/dict/unsendable",
                 ))
             }
         };
-
-        self.flags.push(syn::Expr::Path(path));
         Ok(())
     }
 }
@@ -386,6 +384,16 @@ fn impl_class(
         quote! {}
     };
 
+    let thread_checker = if attr.has_unsendable {
+        quote! { pyo3::pyclass::ThreadCheckerImpl<#cls> }
+    } else if attr.has_extends {
+        quote! {
+            pyo3::pyclass::ThreadCheckerInherited<#cls, <#cls as pyo3::type_object::PyTypeInfo>::BaseType>
+        }
+    } else {
+        quote! { pyo3::pyclass::ThreadCheckerStub<#cls> }
+    };
+
     Ok(quote! {
         unsafe impl pyo3::type_object::PyTypeInfo for #cls {
             type Type = #cls;
@@ -424,6 +432,10 @@ fn impl_class(
             type Target = pyo3::PyRefMut<'a, #cls>;
         }
 
+        unsafe impl pyo3::pyclass::PyClassSend for #cls {
+            type ThreadChecker = #thread_checker;
+        }
+
         #into_pyobject
 
         #impl_inventory
@@ -433,7 +445,6 @@ fn impl_class(
         #extra
 
         #gc_impl
-
     })
 }
 

src/derive_utils.rs

@@ -7,7 +7,7 @@
 use crate::err::{PyErr, PyResult};
 use crate::exceptions::TypeError;
 use crate::instance::PyNativeType;
-use crate::pyclass::PyClass;
+use crate::pyclass::{PyClass, PyClassThreadChecker};
 use crate::types::{PyAny, PyDict, PyModule, PyTuple};
 use crate::{ffi, GILPool, IntoPy, PyCell, Python};
 use std::cell::UnsafeCell;
@@ -162,13 +162,15 @@ pub trait PyBaseTypeUtils {
     type WeakRef;
     type LayoutAsBase;
     type BaseNativeType;
+    type ThreadChecker: PyClassThreadChecker;
 }
 
 impl<T: PyClass> PyBaseTypeUtils for T {
     type Dict = T::Dict;
     type WeakRef = T::WeakRef;
     type LayoutAsBase = crate::pycell::PyCellInner<T>;
     type BaseNativeType = T::BaseNativeType;
+    type ThreadChecker = T::ThreadChecker;
 }
 
 /// Utility trait to enable &PyClass as a pymethod/function argument

src/pycell.rs

@@ -1,10 +1,11 @@
 //! Includes `PyCell` implementation.
 use crate::conversion::{AsPyPointer, FromPyPointer, ToPyObject};
+use crate::pyclass::{PyClass, PyClassThreadChecker};
 use crate::pyclass_init::PyClassInitializer;
 use crate::pyclass_slots::{PyClassDict, PyClassWeakRef};
 use crate::type_object::{PyBorrowFlagLayout, PyLayout, PySizedLayout, PyTypeInfo};
 use crate::types::PyAny;
-use crate::{ffi, FromPy, PyClass, PyErr, PyNativeType, PyObject, PyResult, Python};
+use crate::{ffi, FromPy, PyErr, PyNativeType, PyObject, PyResult, Python};
 use std::cell::{Cell, UnsafeCell};
 use std::fmt;
 use std::mem::ManuallyDrop;
@@ -161,6 +162,7 @@ pub struct PyCell<T: PyClass> {
     inner: PyCellInner<T>,
     dict: T::Dict,
     weakref: T::WeakRef,
+    thread_checker: T::ThreadChecker,
 }
 
 unsafe impl<T: PyClass> PyNativeType for PyCell<T> {}
@@ -227,6 +229,7 @@ impl<T: PyClass> PyCell<T> {
     /// }
     /// ```
     pub fn try_borrow(&self) -> Result<PyRef<'_, T>, PyBorrowError> {
+        self.thread_checker.ensure();
         let flag = self.inner.get_borrow_flag();
         if flag == BorrowFlag::HAS_MUTABLE_BORROW {
             Err(PyBorrowError { _private: () })
@@ -258,6 +261,7 @@ impl<T: PyClass> PyCell<T> {
     /// assert!(c.try_borrow_mut().is_ok());
     /// ```
     pub fn try_borrow_mut(&self) -> Result<PyRefMut<'_, T>, PyBorrowMutError> {
+        self.thread_checker.ensure();
         if self.inner.get_borrow_flag() != BorrowFlag::UNUSED {
             Err(PyBorrowMutError { _private: () })
         } else {
@@ -296,6 +300,7 @@ impl<T: PyClass> PyCell<T> {
     /// }
     /// ```
     pub unsafe fn try_borrow_unguarded(&self) -> Result<&T, PyBorrowError> {
+        self.thread_checker.ensure();
         if self.inner.get_borrow_flag() == BorrowFlag::HAS_MUTABLE_BORROW {
             Err(PyBorrowError { _private: () })
         } else {
@@ -352,6 +357,7 @@ impl<T: PyClass> PyCell<T> {
         let self_ = base as *mut Self;
         (*self_).dict = T::Dict::new();
         (*self_).weakref = T::WeakRef::new();
+        (*self_).thread_checker = T::ThreadChecker::new();
         Ok(self_)
     }
 }

src/pyclass.rs

@@ -2,13 +2,15 @@
 use crate::class::methods::{PyClassAttributeDef, PyMethodDefType, PyMethods};
 use crate::class::proto_methods::PyProtoMethods;
 use crate::conversion::{AsPyPointer, FromPyPointer};
+use crate::derive_utils::PyBaseTypeUtils;
 use crate::pyclass_slots::{PyClassDict, PyClassWeakRef};
 use crate::type_object::{type_flags, PyLayout};
 use crate::types::PyAny;
 use crate::{class, ffi, PyCell, PyErr, PyNativeType, PyResult, PyTypeInfo, Python};
 use std::ffi::CString;
+use std::marker::PhantomData;
 use std::os::raw::c_void;
-use std::ptr;
+use std::{ptr, thread};
 
 #[inline]
 pub(crate) unsafe fn default_new<T: PyTypeInfo>(
@@ -91,10 +93,10 @@ pub(crate) unsafe fn tp_free_fallback(obj: *mut ffi::PyObject) {
 pub trait PyClass:
     PyTypeInfo<Layout = PyCell<Self>, AsRefTarget = PyCell<Self>>
     + Sized
+    + PyClassSend
     + PyClassAlloc
     + PyMethods
     + PyProtoMethods
-    + Send
 {
     /// Specify this class has `#[pyclass(dict)]` or not.
     type Dict: PyClassDict;
@@ -308,3 +310,61 @@ fn py_class_properties<T: PyMethods>() -> Vec<ffi::PyGetSetDef> {
 
     defs.values().cloned().collect()
 }
+
+pub unsafe trait PyClassSend {
+    type ThreadChecker: PyClassThreadChecker;
+}
+
+#[doc(hidden)]
+pub trait PyClassThreadChecker: Sized {
+    fn ensure(&self);
+    fn new() -> Self;
+    private_decl! {}
+}
+
+/// Stub checker for `Send` types.
+#[doc(hidden)]
+pub struct ThreadCheckerStub<T: Send>(PhantomData<T>);
+
+impl<T: Send> PyClassThreadChecker for ThreadCheckerStub<T> {
+    fn ensure(&self) {}
+    fn new() -> Self {
+        ThreadCheckerStub(PhantomData)
+    }
+    private_impl! {}
+}
+
+/// Thread checker for unsendable types.
+/// Panics when the value is accessed by another thread.
+#[doc(hidden)]
+pub struct ThreadCheckerImpl<T>(thread::ThreadId, PhantomData<T>);
+
+impl<T> PyClassThreadChecker for ThreadCheckerImpl<T> {
+    fn ensure(&self) {
+        if thread::current().id() != self.0 {
+            panic!(
+                "{} is unsendable, but sent to another thread!",
+                std::any::type_name::<T>()
+            );
+        }
+    }
+    fn new() -> Self {
+        ThreadCheckerImpl(thread::current().id(), PhantomData)
+    }
+    private_impl! {}
+}
+
+/// Thread checker for types that have `Send` and `extends=...`.
+/// Ensures that `T: Send` and the parent is not accessed by another thread.
+#[doc(hidden)]
+pub struct ThreadCheckerInherited<T: Send, U: PyBaseTypeUtils>(PhantomData<T>, U::ThreadChecker);
+
+impl<T: Send, U: PyBaseTypeUtils> PyClassThreadChecker for ThreadCheckerInherited<T, U> {
+    fn ensure(&self) {
+        self.1.ensure();
+    }
+    fn new() -> Self {
+        ThreadCheckerInherited(PhantomData, U::ThreadChecker::new())
+    }
+    private_impl! {}
+}

src/types/mod.rs

@@ -75,6 +75,7 @@ macro_rules! pyobject_native_type {
             type WeakRef = $crate::pyclass_slots::PyClassDummySlot;
             type LayoutAsBase = $crate::pycell::PyCellBase<$name>;
             type BaseNativeType = $name;
+            type ThreadChecker = $crate::pyclass::ThreadCheckerStub<$crate::PyObject>;
         }
         pyobject_native_type_named!($name $(,$type_param)*);
         pyobject_native_type_convert!($name, $layout, $typeobject, $module, $checkfunction $(,$type_param)*);

tests/test_class_basics.rs

@@ -163,3 +163,55 @@ fn class_with_object_field() {
     py_assert!(py, ty, "ty(5).value == 5");
     py_assert!(py, ty, "ty(None).value == None");
 }
+
+#[pyclass(unsendable)]
+struct UnsendableBase {
+    rc: std::rc::Rc<usize>,
+}
+
+#[pymethods]
+impl UnsendableBase {
+    fn value(&self) -> usize {
+        *self.rc.as_ref()
+    }
+}
+
+#[pyclass(extends=UnsendableBase)]
+struct UnsendableChild {}
+
+/// If a class is marked as `unsendable`, it panics when accessed by another thread.
+#[test]
+fn panic_unsendable() {
+    let gil = Python::acquire_gil();
+    let py = gil.python();
+    let base = || UnsendableBase {
+        rc: std::rc::Rc::new(0),
+    };
+    let unsendable_base = PyCell::new(py, base()).unwrap();
+    let unsendable_child = PyCell::new(py, (UnsendableChild {}, base())).unwrap();
+
+    let source = pyo3::indoc::indoc!(
+        r#"
+def value():
+    return unsendable.value()
+
+import concurrent.futures
+executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
+future = executor.submit(value)
+try:
+    result = future.result()
+    assert False, 'future must panic'
+except BaseException as e:
+    assert str(e) == 'test_class_basics::UnsendableBase is unsendable, but sent to another thread!'
+"#
+    );
+    let globals = PyModule::import(py, "__main__").unwrap().dict();
+    let test = |unsendable| {
+        globals.set_item("unsendable", unsendable).unwrap();
+        py.run(source, Some(globals), None)
+            .map_err(|e| e.print(py))
+            .unwrap();
+    };
+    test(unsendable_base.as_ref());
+    test(unsendable_child.as_ref());
+}

tests/ui/invalid_pyclass_args.stderr

@@ -22,7 +22,7 @@ error: Expected string literal (e.g., "my_mod")
 12 | #[pyclass(module = my_module)]
    |                    ^^^^^^^^^
 
-error: Expected one of gc/weakref/subclass/dict
+error: Expected one of gc/weakref/subclass/dict/unsendable
   --> $DIR/invalid_pyclass_args.rs:15:11
    |
 15 | #[pyclass(weakrev)]