Skip to content

Commit 5e6c733

Browse files
IcxoluIcxolu
committed
improve PyCapsule constructors
1 parent d07a1ab commit 5e6c733

5 files changed

Lines changed: 143 additions & 32 deletions

File tree

‎newsfragments/5881.added.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
added `PyCapsule::new_with_value` and `PyCapsule::new_with_value_and_destructor`

‎src/instance.rs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2753,10 +2753,10 @@ a = A()
27532753
method: impl FnOnce(*mut ffi::PyObject) -> Bound<'py, PyAny>,
27542754
) {
27552755
let mut dropped = false;
2756-
let capsule = PyCapsule::new_with_destructor(
2756+
let capsule = PyCapsule::new_with_value_and_destructor(
27572757
py,
27582758
(&mut dropped) as *mut _ as usize,
2759-
None,
2759+
c"bound_from_borrowed_ptr_constructors",
27602760
|ptr, _| unsafe { std::ptr::write(ptr as *mut bool, true) },
27612761
)
27622762
.unwrap();
@@ -2795,10 +2795,10 @@ a = A()
27952795
method: impl FnOnce(&*mut ffi::PyObject) -> Borrowed<'_, 'py, PyAny>,
27962796
) {
27972797
let mut dropped = false;
2798-
let capsule = PyCapsule::new_with_destructor(
2798+
let capsule = PyCapsule::new_with_value_and_destructor(
27992799
py,
28002800
(&mut dropped) as *mut _ as usize,
2801-
None,
2801+
c"borrowed_ptr_constructors",
28022802
|ptr, _| unsafe { std::ptr::write(ptr as *mut bool, true) },
28032803
)
28042804
.unwrap();

‎src/types/capsule.rs‎

Lines changed: 130 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ use std::ptr::{self, NonNull};
3535
///
3636
/// let r = Python::attach(|py| -> PyResult<()> {
3737
/// let foo = Foo { val: 123 };
38-
/// let capsule = PyCapsule::new(py, foo, Some(c"builtins.capsule".to_owned()))?;
38+
/// let capsule = PyCapsule::new_with_value(py, foo, c"builtins.capsule")?;
3939
///
4040
/// let module = PyModule::import(py, "builtins")?;
4141
/// module.add("capsule", capsule)?;
@@ -69,27 +69,76 @@ impl PyCapsule {
6969
/// // this can be c"foo" on Rust 1.77+
7070
/// const NAME: &CStr = c"foo";
7171
///
72+
/// # fn main() -> PyResult<()> {
7273
/// Python::attach(|py| {
73-
/// let capsule = PyCapsule::new(py, 123_u32, Some(NAME.to_owned())).unwrap();
74-
/// let val: NonNull<u32> = capsule.pointer_checked(Some(NAME)).unwrap().cast();
74+
/// let capsule = PyCapsule::new_with_value(py, 123_u32, NAME)?;
75+
/// let val: NonNull<u32> = capsule.pointer_checked(Some(NAME))?.cast();
7576
/// assert_eq!(unsafe { *val.as_ref() }, 123);
76-
/// });
77+
/// # Ok(())
78+
/// })
79+
/// # }
7780
/// ```
7881
///
7982
/// However, attempting to construct a `PyCapsule` with a zero-sized type will not compile:
8083
///
8184
/// ```compile_fail
8285
/// use pyo3::{prelude::*, types::PyCapsule};
8386
///
87+
/// # fn main() -> PyResult<()> {
8488
/// Python::attach(|py| {
85-
/// let capsule = PyCapsule::new(py, (), None).unwrap(); // Oops! `()` is zero sized!
86-
/// });
89+
/// let capsule = PyCapsule::new_with_value(py, (), c"foo")?; // Oops! `()` is zero sized!
90+
/// # Ok(())
91+
/// })
92+
/// # }
8793
/// ```
94+
pub fn new_with_value<'py, T>(
95+
py: Python<'py>,
96+
value: T,
97+
name: &'static CStr,
98+
) -> PyResult<Bound<'py, Self>>
99+
where
100+
T: 'static + Send + AssertNotZeroSized,
101+
{
102+
AssertNotZeroSized::assert_not_zero_sized(&value);
103+
104+
// NOTE: Not implemented in terms of `new_with_value_and_destructor` as this allows for
105+
// storing the `Box<T>` directly without allocating additionally for the destructor
106+
//
107+
// SAFETY: `Box` pointers are guaranteed to be non-null
108+
let val = unsafe { NonNull::new_unchecked(Box::into_raw(Box::new(value))) };
109+
110+
unsafe extern "C" fn destructor<T>(capsule: *mut ffi::PyObject) {
111+
// SAFETY: `capsule` is known to be a borrowed reference to the capsule being destroyed
112+
let name = unsafe { ffi::PyCapsule_GetName(capsule) };
113+
114+
// SAFETY:
115+
// - `capsule` is known to be a borrowed reference to the capsule being destroyed
116+
// - `name` is known to be the capsule's name
117+
let ptr = unsafe { ffi::PyCapsule_GetPointer(capsule, name) };
118+
119+
// SAFETY: `capsule` was knowingly constructed from a `Box<T>` and is now being
120+
// destroyed, so we reconstruct the Box and drop it.
121+
let _ = unsafe { Box::<T>::from_raw(ptr.cast()) };
122+
}
123+
124+
// SAFETY:
125+
// - `val` is aligned and valid for a `T` until the destructor runs
126+
// - `destructor` will deallocate the `Box`
127+
// - `destructor` can be called from any thread as `T: Send`
128+
// - `destructor` does not panic
129+
unsafe {
130+
Self::new_with_pointer_and_destructor(py, val.cast(), name, Some(destructor::<T>))
131+
}
132+
}
133+
134+
/// See [`PyCapsule::new_with_value`]
135+
#[deprecated(since = "0.29.0", note = "use `PyCapsule::new_with_value` instead")]
88136
pub fn new<T: 'static + Send + AssertNotZeroSized>(
89137
py: Python<'_>,
90138
value: T,
91139
name: Option<CString>,
92140
) -> PyResult<Bound<'_, Self>> {
141+
#[allow(deprecated)]
93142
Self::new_with_destructor(py, value, name, |_, _| {})
94143
}
95144

@@ -100,6 +149,54 @@ impl PyCapsule {
100149
///
101150
/// The `destructor` must be `Send`, because there is no guarantee which thread it will eventually
102151
/// be called from.
152+
///
153+
/// # Note
154+
/// If `destructor` panics the process will (safely) abort
155+
pub fn new_with_value_and_destructor<'py, T, F>(
156+
py: Python<'py>,
157+
value: T,
158+
name: &'static CStr,
159+
destructor: F,
160+
) -> PyResult<Bound<'py, Self>>
161+
where
162+
T: 'static + Send + AssertNotZeroSized,
163+
F: FnOnce(T, *mut c_void) + Send,
164+
{
165+
AssertNotZeroSized::assert_not_zero_sized(&value);
166+
167+
// Sanity check for capsule layout
168+
debug_assert_eq!(offset_of!(CapsuleContents::<T, F>, value), 0);
169+
170+
// SAFETY: `Box` pointers are guaranteed to be non-null
171+
let val = unsafe {
172+
NonNull::new_unchecked(Box::into_raw(Box::new(CapsuleContents {
173+
value,
174+
destructor,
175+
name: None,
176+
})))
177+
};
178+
179+
// SAFETY:
180+
// - `val` is an aligned and valid pointer to a `CapsuleContents` until the destructor runs
181+
// - `CapsuleContents` is `#[repr(C)]` with `T` as its first field, so `val` may be used as
182+
// an aligned and valid pointer to a `T`
183+
// - `capsule_destructor` will deallocate the `Box` and call the user provided `destructor`
184+
// - `capsule_destructor` can be called from any thread as `T: Send` and `F: Send`
185+
unsafe {
186+
Self::new_with_pointer_and_destructor(
187+
py,
188+
val.cast(),
189+
name,
190+
Some(capsule_destructor::<T, F>),
191+
)
192+
}
193+
}
194+
195+
/// See [`PyCapsule::new_with_value_and_destructor`]
196+
#[deprecated(
197+
since = "0.29.0",
198+
note = "use `PyCapsule::new_with_value_and_destructor` instead"
199+
)]
103200
pub fn new_with_destructor<
104201
T: 'static + Send + AssertNotZeroSized,
105202
F: FnOnce(T, *mut c_void) + Send,
@@ -307,6 +404,7 @@ pub trait PyCapsuleMethods<'py>: crate::sealed::Sealed {
307404
/// use std::sync::mpsc::{channel, Sender};
308405
/// use pyo3::{prelude::*, types::PyCapsule};
309406
///
407+
/// # fn main() -> PyResult<()> {
310408
/// let (tx, rx) = channel::<String>();
311409
///
312410
/// fn destructor(val: u32, context: *mut c_void) {
@@ -315,15 +413,21 @@ pub trait PyCapsuleMethods<'py>: crate::sealed::Sealed {
315413
/// }
316414
///
317415
/// Python::attach(|py| {
318-
/// let capsule =
319-
/// PyCapsule::new_with_destructor(py, 123, None, destructor as fn(u32, *mut c_void))
320-
/// .unwrap();
416+
/// let capsule = PyCapsule::new_with_value_and_destructor(
417+
/// py,
418+
/// 123,
419+
/// c"foo",
420+
/// destructor as fn(u32, *mut c_void)
421+
/// )?;
321422
/// let context = Box::new(tx); // `Sender<String>` is our context, box it up and ship it!
322-
/// capsule.set_context(Box::into_raw(context).cast()).unwrap();
423+
/// capsule.set_context(Box::into_raw(context).cast())?;
323424
/// // This scope will end, causing our destructor to be called...
324-
/// });
425+
/// # Ok::<_, PyErr>(())
426+
/// })?;
325427
///
326428
/// assert_eq!(rx.recv(), Ok("Destructor called!".to_string()));
429+
/// # Ok(())
430+
/// }
327431
/// ```
328432
fn set_context(&self, context: *mut c_void) -> PyResult<()>;
329433

@@ -524,6 +628,7 @@ struct CapsuleContents<T: 'static + Send, D: FnOnce(T, *mut c_void) + Send> {
524628
/// Destructor to be used by the capsule
525629
destructor: D,
526630
/// Name used when creating the capsule
631+
// TODO: remove this field when `PyCapsule::new` and `PyCapsule::new_with_destructor` are removed
527632
name: Option<CString>,
528633
}
529634

@@ -634,7 +739,7 @@ mod tests {
634739
Python::attach(|py| {
635740
let foo = Foo { val: 123 };
636741

637-
let cap = PyCapsule::new(py, foo, Some(NAME.to_owned())).unwrap();
742+
let cap = PyCapsule::new_with_value(py, foo, NAME).unwrap();
638743
assert!(cap.is_valid_checked(Some(NAME)));
639744

640745
let foo_capi = cap.pointer_checked(Some(NAME)).unwrap().cast::<Foo>();
@@ -659,7 +764,7 @@ mod tests {
659764
}
660765

661766
let cap: Py<PyCapsule> = Python::attach(|py| {
662-
let cap = PyCapsule::new(py, foo as fn(u32) -> u32, Some(NAME.to_owned())).unwrap();
767+
let cap = PyCapsule::new_with_value(py, foo as fn(u32) -> u32, NAME).unwrap();
663768
cap.into()
664769
});
665770

@@ -677,7 +782,7 @@ mod tests {
677782
#[test]
678783
fn test_pycapsule_context() {
679784
Python::attach(|py| {
680-
let cap = PyCapsule::new(py, 0, Some(NAME.to_owned())).unwrap();
785+
let cap = PyCapsule::new_with_value(py, 0, NAME).unwrap();
681786

682787
let c = cap.context().unwrap();
683788
assert!(c.is_null());
@@ -703,7 +808,7 @@ mod tests {
703808
let foo = Foo { val: 123 };
704809
let name = c"builtins.capsule";
705810

706-
let capsule = PyCapsule::new(py, foo, Some(name.to_owned())).unwrap();
811+
let capsule = PyCapsule::new_with_value(py, foo, name).unwrap();
707812

708813
let module = PyModule::import(py, "builtins").unwrap();
709814
module.add("capsule", capsule).unwrap();
@@ -724,7 +829,7 @@ mod tests {
724829
fn test_vec_storage() {
725830
let cap: Py<PyCapsule> = Python::attach(|py| {
726831
let stuff: Vec<u8> = vec![1, 2, 3, 4];
727-
let cap = PyCapsule::new(py, stuff, Some(NAME.to_owned())).unwrap();
832+
let cap = PyCapsule::new_with_value(py, stuff, NAME).unwrap();
728833
cap.into()
729834
});
730835

@@ -744,7 +849,7 @@ mod tests {
744849
let context: Vec<u8> = vec![1, 2, 3, 4];
745850

746851
let cap: Py<PyCapsule> = Python::attach(|py| {
747-
let cap = PyCapsule::new(py, 0, Some(NAME.to_owned())).unwrap();
852+
let cap = PyCapsule::new_with_value(py, 0, NAME).unwrap();
748853
cap.set_context(Box::into_raw(Box::new(&context)).cast())
749854
.unwrap();
750855

@@ -771,8 +876,7 @@ mod tests {
771876
}
772877

773878
Python::attach(move |py| {
774-
let cap =
775-
PyCapsule::new_with_destructor(py, 0, Some(NAME.to_owned()), destructor).unwrap();
879+
let cap = PyCapsule::new_with_value_and_destructor(py, 0, NAME, destructor).unwrap();
776880
cap.set_context(Box::into_raw(Box::new(tx)).cast()).unwrap();
777881
});
778882

@@ -781,6 +885,7 @@ mod tests {
781885
}
782886

783887
#[test]
888+
#[allow(deprecated)]
784889
fn test_pycapsule_no_name() {
785890
Python::attach(|py| {
786891
let cap = PyCapsule::new(py, 0usize, None).unwrap();
@@ -869,7 +974,7 @@ mod tests {
869974
#[test]
870975
fn test_pycapsule_pointer_checked_wrong_name() {
871976
Python::attach(|py| {
872-
let cap = PyCapsule::new(py, 123u32, Some(c"correct.name".to_owned())).unwrap();
977+
let cap = PyCapsule::new_with_value(py, 123u32, c"correct.name").unwrap();
873978

874979
// Requesting with wrong name should fail
875980
let result = cap.pointer_checked(Some(c"wrong.name"));
@@ -882,6 +987,7 @@ mod tests {
882987
}
883988

884989
#[test]
990+
#[allow(deprecated)]
885991
fn test_pycapsule_pointer_checked_none_vs_some() {
886992
Python::attach(|py| {
887993
// Capsule with no name
@@ -899,7 +1005,7 @@ mod tests {
8991005
#[test]
9001006
fn test_pycapsule_is_valid_checked_wrong_name() {
9011007
Python::attach(|py| {
902-
let cap = PyCapsule::new(py, 123u32, Some(c"correct.name".to_owned())).unwrap();
1008+
let cap = PyCapsule::new_with_value(py, 123u32, c"correct.name").unwrap();
9031009

9041010
// Should be valid with correct name
9051011
assert!(cap.is_valid_checked(Some(c"correct.name")));
@@ -913,6 +1019,7 @@ mod tests {
9131019
}
9141020

9151021
#[test]
1022+
#[allow(deprecated)]
9161023
fn test_pycapsule_is_valid_checked_no_name() {
9171024
Python::attach(|py| {
9181025
let cap = PyCapsule::new(py, 123u32, None).unwrap();
@@ -928,7 +1035,7 @@ mod tests {
9281035
#[test]
9291036
fn test_pycapsule_context_on_invalid_capsule() {
9301037
Python::attach(|py| {
931-
let cap = PyCapsule::new(py, 123u32, Some(NAME.to_owned())).unwrap();
1038+
let cap = PyCapsule::new_with_value(py, 123u32, NAME).unwrap();
9321039

9331040
// Invalidate the capsule
9341041
// SAFETY: intentionally breaking the capsule for testing
@@ -957,7 +1064,7 @@ mod tests {
9571064
fn test_pycapsule_import_wrong_attribute() {
9581065
Python::attach(|py| {
9591066
// Create a capsule and register it
960-
let cap = PyCapsule::new(py, 123u32, Some(c"builtins.test_cap".to_owned())).unwrap();
1067+
let cap = PyCapsule::new_with_value(py, 123u32, c"builtins.test_cap").unwrap();
9611068
let module = crate::prelude::PyModule::import(py, "builtins").unwrap();
9621069
module.add("test_cap", cap).unwrap();
9631070

‎src/types/function.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -91,13 +91,13 @@ impl PyCFunction {
9191
pymethods::PyMethodDef::cfunction_with_keywords(name, run_closure::<F, R>, doc);
9292
let def = method_def.into_raw();
9393

94-
let capsule = PyCapsule::new(
94+
let capsule = PyCapsule::new_with_value(
9595
py,
9696
ClosureDestructor::<F> {
9797
closure,
9898
def: UnsafeCell::new(def),
9999
},
100-
Some(CLOSURE_CAPSULE_NAME.to_owned()),
100+
CLOSURE_CAPSULE_NAME,
101101
)?;
102102

103103
let data: NonNull<ClosureDestructor<F>> =

‎tests/test_inheritance.rs‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -253,9 +253,12 @@ mod inheriting_native_type {
253253
Python::attach(|py| {
254254
let dropped = Arc::new(AtomicBool::new(false));
255255
let destructor_drop = Arc::clone(&dropped);
256-
let item = PyCapsule::new_with_destructor(py, 0, None, move |_, _| {
257-
destructor_drop.store(true, Ordering::Relaxed)
258-
})
256+
let item = PyCapsule::new_with_value_and_destructor(
257+
py,
258+
0,
259+
c"inherit_dict_drop",
260+
move |_, _| destructor_drop.store(true, Ordering::Relaxed),
261+
)
259262
.unwrap();
260263

261264
let dict_sub = pyo3::Py::new(py, DictWithName::new()).unwrap();

0 commit comments

Comments
 (0)