#![deny(clippy::undocumented_unsafe_blocks)]
use crate::{PyAny, ffi};
#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))]
use crate::{Borrowed, PyResult, Python, err};
#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))]
use core::ffi::c_int;
#[repr(transparent)]
pub struct PyContext(PyAny);
pyobject_native_type_core!(
PyContext,
pyobject_native_static_type_object!(ffi::PyContext_Type),
"contextvars",
"Context",
#module=Some("contextvars"),
#checkfunction=ffi::PyContext_CheckExact
);
#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))]
impl PyContext {
#[doc(alias = "PyContext_AddWatcher")]
pub fn add_watcher(
py: Python<'_>,
callback: WatchCallback,
) -> PyResult<BoundContextWatcherGuard<'_>> {
let watcher_id = err::error_on_negative_with_result(py, unsafe {
ffi::PyContext_AddWatcher(callback.0)
})?;
Ok(BoundContextWatcherGuard {
watcher_id,
py,
active: true,
})
}
}
#[doc(alias = "PyContextEvent")]
#[derive(Debug)]
#[non_exhaustive]
#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))]
pub enum ContextEvent<'a, 'py> {
Switched(Option<Borrowed<'a, 'py, PyContext>>),
Unknown {
raw_event: ffi::PyContextEvent,
object: Option<Borrowed<'a, 'py, PyAny>>,
},
}
#[must_use = "dropping the guard immediately unregisters the context watcher"]
#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))]
pub struct BoundContextWatcherGuard<'py> {
watcher_id: c_int,
py: Python<'py>,
active: bool,
}
#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))]
impl BoundContextWatcherGuard<'_> {
#[doc(alias = "PyContext_ClearWatcher")]
pub fn clear(mut self) -> PyResult<()> {
self.active = false;
clear_watcher(self.py, self.watcher_id)
}
pub fn unbind(mut self) -> ContextWatcherGuard {
self.active = false;
ContextWatcherGuard {
watcher_id: self.watcher_id,
active: true,
}
}
}
#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))]
impl Drop for BoundContextWatcherGuard<'_> {
fn drop(&mut self) {
if !self.active {
return;
}
self.active = false;
clear_watcher_on_drop(self.py, self.watcher_id);
}
}
#[must_use = "dropping the guard unregisters the context watcher"]
#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))]
pub struct ContextWatcherGuard {
watcher_id: c_int,
active: bool,
}
#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))]
impl ContextWatcherGuard {
#[doc(alias = "PyContext_ClearWatcher")]
pub fn clear(mut self, py: Python<'_>) -> PyResult<()> {
self.active = false;
clear_watcher(py, self.watcher_id)
}
pub fn into_bound<'py>(mut self, py: Python<'py>) -> BoundContextWatcherGuard<'py> {
self.active = false;
BoundContextWatcherGuard {
watcher_id: self.watcher_id,
py,
active: true,
}
}
}
#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))]
impl Drop for ContextWatcherGuard {
fn drop(&mut self) {
if !self.active {
return;
}
self.active = false;
let watcher_id = self.watcher_id;
let _ = Python::try_attach(|py| clear_watcher_on_drop(py, watcher_id));
}
}
#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))]
fn clear_watcher(py: Python<'_>, watcher_id: c_int) -> PyResult<()> {
err::error_on_negative(py, unsafe { ffi::PyContext_ClearWatcher(watcher_id) })
}
#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))]
fn clear_watcher_on_drop(_py: Python<'_>, watcher_id: c_int) {
unsafe {
let pending_exception = ffi::PyErr_GetRaisedException();
if ffi::PyContext_ClearWatcher(watcher_id) == -1 {
ffi::PyErr_WriteUnraisable(core::ptr::null_mut());
}
if !pending_exception.is_null() {
ffi::PyErr_Clear();
ffi::PyErr_SetRaisedException(pending_exception);
}
}
}
#[repr(transparent)]
#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))]
pub struct WatchCallback(ffi::PyContext_WatchCallback);
#[macro_export]
#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))]
macro_rules! watch_callback {
($callback:path) => {{
struct Callback;
impl $crate::types::context::impl_::ContextWatcherCallbackDef for Callback {
const CALLBACK: $crate::types::context::impl_::ContextWatcherCallback = $callback;
}
$crate::types::context::impl_::new_watch_callback::<Callback>()
}};
}
#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))]
pub use crate::watch_callback;
#[doc(hidden)]
#[cfg(all(Py_3_14, not(Py_GIL_DISABLED)))]
pub mod impl_ {
use crate::{ffi_ptr_ext::FfiPtrExt, types::PyAnyMethods};
use super::*;
pub type ContextWatcherCallback =
for<'a, 'py> fn(Python<'py>, ContextEvent<'a, 'py>) -> PyResult<()>;
pub trait ContextWatcherCallbackDef {
const CALLBACK: ContextWatcherCallback;
}
pub fn new_watch_callback<Callback: ContextWatcherCallbackDef>() -> WatchCallback {
WatchCallback(context_watcher::<Callback>)
}
unsafe fn event_from_raw<'a, 'py>(
py: Python<'py>,
event: ffi::PyContextEvent,
object: *mut ffi::PyObject,
) -> PyResult<ContextEvent<'a, 'py>> {
match event {
ffi::Py_CONTEXT_SWITCHED => {
let object = unsafe { object.assume_borrowed_unchecked(py) };
if object.is_none() {
Ok(ContextEvent::Switched(None))
} else {
Ok(ContextEvent::Switched(Some(object.cast()?)))
}
}
raw_event => {
let object = unsafe { object.assume_borrowed_or_opt(py) };
Ok(ContextEvent::Unknown { raw_event, object })
}
}
}
pub unsafe extern "C" fn context_watcher<Callback: ContextWatcherCallbackDef>(
event: ffi::PyContextEvent,
object: *mut ffi::PyObject,
) -> c_int {
let pending_exception = unsafe { ffi::PyErr_GetRaisedException() };
let result = unsafe {
crate::impl_::trampoline::trampoline(|py| {
let event = event_from_raw(py, event, object)?;
(Callback::CALLBACK)(py, event)?;
if crate::PyErr::occurred(py) {
return Err(crate::PyErr::fetch(py));
}
Ok(0)
})
};
if pending_exception.is_null() {
return result;
}
unsafe {
if result == -1 {
ffi::PyErr_WriteUnraisable(object);
}
ffi::PyErr_Clear();
ffi::PyErr_SetRaisedException(pending_exception);
}
0
}
}
#[cfg(all(test, Py_3_14, not(Py_GIL_DISABLED)))]
mod watcher_tests {
use super::impl_::{ContextWatcherCallback, ContextWatcherCallbackDef, context_watcher};
use super::{ContextEvent, PyContext};
use crate::exceptions::{PyRuntimeError, PyValueError};
use crate::platform::sync::nonpoison::{Mutex, MutexGuard};
#[cfg(feature = "macros")]
use crate::test_utils::UnraisableCapture;
use crate::types::PyAnyMethods;
use crate::{PyErr, PyResult, Python, ffi};
use alloc::string::ToString;
use core::sync::atomic::{AtomicBool, AtomicU32, AtomicUsize, Ordering};
use static_assertions::{assert_impl_all, assert_not_impl_any};
static WATCHER_TEST_MUTEX: Mutex<()> = Mutex::new(());
static SWITCH_COUNT: AtomicUsize = AtomicUsize::new(0);
static SAW_CONTEXT: AtomicBool = AtomicBool::new(false);
fn acquire_watcher_test_lock() -> MutexGuard<'static, ()> {
WATCHER_TEST_MUTEX.lock()
}
fn run_context_switch(py: Python<'_>) {
py.run(
c"import contextvars; contextvars.Context().run(lambda: None)",
None,
None,
)
.unwrap();
}
fn assert_no_context_switches(py: Python<'_>, count_before: usize) {
run_context_switch(py);
assert_eq!(SWITCH_COUNT.load(Ordering::Relaxed), count_before);
}
#[allow(clippy::unnecessary_wraps, reason = "context watcher callback")]
fn record_switch(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> {
if let ContextEvent::Switched(context) = event {
SWITCH_COUNT.fetch_add(1, Ordering::Relaxed);
if let Some(context) = context {
assert!(context.is_exact_instance_of::<PyContext>());
SAW_CONTEXT.store(true, Ordering::Relaxed);
}
}
Ok(())
}
#[test]
fn watcher_is_cleared_on_drop() {
let _guard = acquire_watcher_test_lock();
Python::attach(|py| {
SWITCH_COUNT.store(0, Ordering::Relaxed);
SAW_CONTEXT.store(false, Ordering::Relaxed);
let watcher = PyContext::add_watcher(py, watch_callback!(record_switch)).unwrap();
run_context_switch(py);
let count_after_first_run = SWITCH_COUNT.load(Ordering::Relaxed);
assert!(count_after_first_run >= 2);
assert!(SAW_CONTEXT.load(Ordering::Relaxed));
drop(watcher);
assert_no_context_switches(py, count_after_first_run);
});
}
#[test]
fn watcher_can_be_cleared_explicitly() {
let _guard = acquire_watcher_test_lock();
Python::attach(|py| {
SWITCH_COUNT.store(0, Ordering::Relaxed);
let watcher = PyContext::add_watcher(py, watch_callback!(record_switch)).unwrap();
watcher.clear().unwrap();
assert_no_context_switches(py, 0);
});
}
#[test]
fn multiple_watchers_can_register_the_same_callback() {
let _guard = acquire_watcher_test_lock();
Python::attach(|py| {
SWITCH_COUNT.store(0, Ordering::Relaxed);
let first = PyContext::add_watcher(py, watch_callback!(record_switch)).unwrap();
let second = PyContext::add_watcher(py, watch_callback!(record_switch)).unwrap();
run_context_switch(py);
assert!(SWITCH_COUNT.load(Ordering::Relaxed) >= 4);
drop(first);
let count_with_both = SWITCH_COUNT.load(Ordering::Relaxed);
run_context_switch(py);
assert!(SWITCH_COUNT.load(Ordering::Relaxed) >= count_with_both + 2);
drop(second);
let count_after_drop = SWITCH_COUNT.load(Ordering::Relaxed);
assert_no_context_switches(py, count_after_drop);
});
}
#[test]
fn dropping_watcher_preserves_a_pending_exception() {
let _guard = acquire_watcher_test_lock();
Python::attach(|py| {
let watcher = PyContext::add_watcher(py, watch_callback!(record_switch)).unwrap();
PyValueError::new_err("original error").restore(py);
drop(watcher);
let error = PyErr::fetch(py);
assert!(error.is_instance_of::<PyValueError>(py));
});
}
fn fail_callback(_py: Python<'_>, _event: ContextEvent<'_, '_>) -> PyResult<()> {
Err(PyRuntimeError::new_err("watcher failed"))
}
struct FailingCallback;
impl ContextWatcherCallbackDef for FailingCallback {
const CALLBACK: ContextWatcherCallback = fail_callback;
}
#[test]
fn callback_error_is_returned_without_a_pending_exception() {
Python::attach(|py| {
let result = unsafe {
context_watcher::<FailingCallback>(ffi::Py_CONTEXT_SWITCHED, ffi::Py_None())
};
assert_eq!(result, -1);
let error = PyErr::fetch(py);
assert!(error.is_instance_of::<PyRuntimeError>(py));
});
}
#[allow(clippy::unnecessary_wraps, reason = "context watcher callback")]
fn restore_error_callback(py: Python<'_>, _event: ContextEvent<'_, '_>) -> PyResult<()> {
PyRuntimeError::new_err("watcher restored error").restore(py);
Ok(())
}
struct RestoringErrorCallback;
impl ContextWatcherCallbackDef for RestoringErrorCallback {
const CALLBACK: ContextWatcherCallback = restore_error_callback;
}
#[test]
fn callback_cannot_return_success_with_an_exception_set() {
Python::attach(|py| {
let result = unsafe {
context_watcher::<RestoringErrorCallback>(ffi::Py_CONTEXT_SWITCHED, ffi::Py_None())
};
assert_eq!(result, -1);
let error = PyErr::fetch(py);
assert!(error.is_instance_of::<PyRuntimeError>(py));
assert_eq!(error.to_string(), "RuntimeError: watcher restored error");
});
}
#[test]
#[cfg(feature = "macros")]
fn callback_error_preserves_a_pending_exception() {
Python::attach(|py| {
UnraisableCapture::enter(py, |capture| {
PyValueError::new_err("original error").restore(py);
let result = unsafe {
context_watcher::<FailingCallback>(ffi::Py_CONTEXT_SWITCHED, ffi::Py_None())
};
assert_eq!(result, 0);
let original_error = PyErr::fetch(py);
assert!(original_error.is_instance_of::<PyValueError>(py));
assert_eq!(original_error.to_string(), "ValueError: original error");
let (watcher_error, object) =
capture.take_capture().expect("missing unraisable error");
assert!(watcher_error.is_instance_of::<PyRuntimeError>(py));
assert!(object.is_none());
});
});
}
#[test]
#[cfg(feature = "macros")]
fn registered_callback_errors_are_unraisable() {
let _guard = acquire_watcher_test_lock();
Python::attach(|py| {
UnraisableCapture::enter(py, |capture| {
let watcher = PyContext::add_watcher(py, watch_callback!(fail_callback)).unwrap();
run_context_switch(py);
let (watcher_error, _) = capture.take_capture().expect("missing unraisable error");
assert!(watcher_error.is_instance_of::<PyRuntimeError>(py));
drop(watcher);
});
});
}
#[cfg(all(wip_feature_std, panic = "unwind"))]
fn panic_callback(_py: Python<'_>, _event: ContextEvent<'_, '_>) -> PyResult<()> {
panic!("context watcher panic")
}
#[cfg(all(wip_feature_std, panic = "unwind"))]
struct PanickingCallback;
#[cfg(all(wip_feature_std, panic = "unwind"))]
impl ContextWatcherCallbackDef for PanickingCallback {
const CALLBACK: ContextWatcherCallback = panic_callback;
}
#[cfg(all(wip_feature_std, panic = "unwind"))]
#[test]
fn callback_panic_does_not_cross_ffi_boundary() {
Python::attach(|py| {
let result = unsafe {
context_watcher::<PanickingCallback>(ffi::Py_CONTEXT_SWITCHED, ffi::Py_None())
};
assert_eq!(result, -1);
assert!(PyErr::occurred(py));
unsafe { ffi::PyErr_Clear() };
});
}
static UNKNOWN_EVENT: AtomicU32 = AtomicU32::new(0);
#[allow(clippy::unnecessary_wraps, reason = "context watcher callback")]
fn record_unknown(_py: Python<'_>, event: ContextEvent<'_, '_>) -> PyResult<()> {
if let ContextEvent::Unknown { raw_event, object } = event {
UNKNOWN_EVENT.store(raw_event, Ordering::Relaxed);
assert!(object.is_none());
}
Ok(())
}
struct UnknownCallback;
impl ContextWatcherCallbackDef for UnknownCallback {
const CALLBACK: ContextWatcherCallback = record_unknown;
}
#[test]
fn unknown_events_are_forwarded() {
const FUTURE_EVENT: ffi::PyContextEvent = 123;
Python::attach(|_py| {
UNKNOWN_EVENT.store(0, Ordering::Relaxed);
let result =
unsafe { context_watcher::<UnknownCallback>(FUTURE_EVENT, core::ptr::null_mut()) };
assert_eq!(result, 0);
assert_eq!(UNKNOWN_EVENT.load(Ordering::Relaxed), FUTURE_EVENT);
});
}
#[test]
fn context_watcher_guard_traits() {
assert_not_impl_any!(super::BoundContextWatcherGuard<'_>: Send, Sync);
assert_impl_all!(super::ContextWatcherGuard: Send, Sync);
}
#[cfg(not(target_arch = "wasm32"))] #[test]
fn unbound_watcher_attaches_on_drop_from_another_thread() {
let _guard = acquire_watcher_test_lock();
SWITCH_COUNT.store(0, Ordering::Relaxed);
let watcher = Python::attach(|py| {
PyContext::add_watcher(py, watch_callback!(record_switch))
.unwrap()
.unbind()
});
Python::attach(run_context_switch);
assert!(SWITCH_COUNT.load(Ordering::Relaxed) >= 2);
std::thread::spawn(move || drop(watcher)).join().unwrap();
let count_after_drop = SWITCH_COUNT.load(Ordering::Relaxed);
Python::attach(|py| assert_no_context_switches(py, count_after_drop));
}
#[test]
fn unbound_watcher_can_be_cleared_with_an_attachment() {
let _guard = acquire_watcher_test_lock();
SWITCH_COUNT.store(0, Ordering::Relaxed);
let watcher = Python::attach(|py| {
PyContext::add_watcher(py, watch_callback!(record_switch))
.unwrap()
.unbind()
});
let count_after_clear = Python::attach(|py| {
watcher.clear(py).unwrap();
SWITCH_COUNT.load(Ordering::Relaxed)
});
Python::attach(|py| assert_no_context_switches(py, count_after_clear));
}
#[test]
fn unbound_watcher_can_be_rebound() {
let _guard = acquire_watcher_test_lock();
SWITCH_COUNT.store(0, Ordering::Relaxed);
let watcher = Python::attach(|py| {
PyContext::add_watcher(py, watch_callback!(record_switch))
.unwrap()
.unbind()
});
Python::attach(|py| {
let watcher = watcher.into_bound(py);
run_context_switch(py);
let count_before_drop = SWITCH_COUNT.load(Ordering::Relaxed);
assert!(count_before_drop >= 2);
drop(watcher);
assert_no_context_switches(py, count_before_drop);
});
}
}
#[cfg(test)]
mod tests {
use super::PyContext;
use crate::Python;
use crate::types::PyAnyMethods;
#[test]
fn context_type() {
Python::attach(|py| {
let context = py
.import(c"contextvars")
.unwrap()
.getattr(c"Context")
.unwrap()
.call0()
.unwrap();
assert!(context.is_exact_instance_of::<PyContext>());
context.cast::<PyContext>().unwrap();
});
}
}