Skip to content

Commit 7e2ef18

Browse files
MatthieuDartiailhCopilotdavidhewitt
authored
Avoid type checks in methods where CPython already guarantees the received type (#5930)
* speicialize self conversion for descriptor slots for which CPython does ensure a bad type cannot be passed from Python * Implement trusted self conversion for all extension-type method and slot wrappers * tests: bless outputs * add a news fragment * address review comments * proper news fragment * wip fixing tests * attempt to fix test * fix broken tests * address review comment * test: bless ui test * test: disable 2 tests based on accessing __getattr__ from type object * tests: broader type ignore in tests testing with known wrong types * test: bless ui tests * fix clippy warning * test: ui fix indentation * test; fix again bad indent in ui test * test: do not test for silly calls in pure Python * test: ui fix formatting * address review comments * tests: attempt to fix ui tests * tests: attempt to fix ui tests * Revert "tests: attempt to fix ui tests" This reverts commit 28d5326. * test: ui attempt to fix tests * fix bad indent * revert bad changes to invalid_pyfunction_argument.default.stderr * label SAFETY comments * disable trusted optimization on PyPy * fixup richcmp on PyPy * fix msrv build * fix clippy --------- Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com> Co-authored-by: David Hewitt <mail@davidhewitt.dev>
1 parent f930199 commit 7e2ef18

15 files changed

Lines changed: 621 additions & 165 deletions

‎newsfragments/5930.changed.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Remove redundant type checks for methods where CPython guarantees the type of `self`

‎pyo3-macros-backend/src/method.rs‎

Lines changed: 123 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -263,6 +263,7 @@ impl FnType {
263263
&self,
264264
cls: Option<&syn::Type>,
265265
error_mode: ExtractErrorMode,
266+
self_conversion: SelfConversionPolicy,
266267
holders: &mut Holders,
267268
ctx: &Ctx,
268269
) -> Option<TokenStream> {
@@ -272,6 +273,7 @@ impl FnType {
272273
Some(st.receiver(
273274
cls.expect("no class given for Fn with a \"self\" receiver"),
274275
error_mode,
276+
self_conversion,
275277
holders,
276278
ctx,
277279
))
@@ -320,6 +322,56 @@ pub enum SelfType {
320322
},
321323
}
322324

325+
#[derive(Clone, Copy, Debug)]
326+
enum SelfConversionPolicyInner {
327+
/// The receiver's type is guaranteed by CPython's slot/method dispatch contract.
328+
/// Used for all extension-type method and slot entrypoints.
329+
Trusted,
330+
/// The receiver's type is verified at runtime. Used for number-protocol
331+
/// binary operator fragments where the CPython dispatch contract does not
332+
/// guarantee the receiver type.
333+
Checked,
334+
}
335+
336+
/// Receiver conversion policy for extension-type method wrappers.
337+
///
338+
/// Controls whether the `self` receiver is validated with a runtime type check
339+
/// (`Checked`) or treated as trusted and cast directly without checking
340+
/// (`Trusted`).
341+
///
342+
/// # Invariant
343+
///
344+
/// The `Trusted` path is valid due to CPython's slot/method receiver contract:
345+
/// when CPython dispatches a method call on an extension type — whether through
346+
/// a type slot or through `tp_methods` — the receiver is guaranteed to be an
347+
/// instance of the owning type (or a compatible subtype). For `tp_methods`
348+
/// entries, CPython's method-wrapper descriptor enforces this before the C
349+
/// function is reached.
350+
///
351+
/// `Checked` should be used in cases where that guarantee does not hold:
352+
/// - Number-protocol binary operator fragments (`__add__`, `__radd__`, …,
353+
/// `__pow__`, `__rpow__`): CPython combines the forward and reflected
354+
/// fragments into a single `nb_add`/`nb_power` slot, and the runtime helper
355+
/// may call the reflected fragment with the operands swapped, meaning `_slf`
356+
/// can arrive with a non-class type. The existing
357+
/// `ExtractErrorMode::NotImplemented` behavior on type mismatch is preserved
358+
/// by using `Checked` for those fragments.
359+
#[derive(Clone, Copy, Debug)]
360+
pub struct SelfConversionPolicy(SelfConversionPolicyInner);
361+
362+
impl SelfConversionPolicy {
363+
pub const fn checked() -> Self {
364+
Self(SelfConversionPolicyInner::Checked)
365+
}
366+
367+
// Using the trusted conversion incorrectly can lead to incorrect runtime
368+
// behavior and memory safety issues, so this is marked `unsafe` and usage
369+
// should be justified by the caller.
370+
pub const unsafe fn trusted() -> Self {
371+
Self(SelfConversionPolicyInner::Trusted)
372+
}
373+
}
374+
323375
#[derive(Clone, Copy)]
324376
pub enum ExtractErrorMode {
325377
NotImplemented,
@@ -346,6 +398,7 @@ impl SelfType {
346398
&self,
347399
cls: &syn::Type,
348400
error_mode: ExtractErrorMode,
401+
self_conversion: SelfConversionPolicy,
349402
holders: &mut Holders,
350403
ctx: &Ctx,
351404
) -> TokenStream {
@@ -367,22 +420,45 @@ impl SelfType {
367420
};
368421
let arg =
369422
quote! { unsafe { #pyo3_path::impl_::extract_argument::#cast_fn(#py, #slf) } };
370-
let method = if *mutable {
371-
syn::Ident::new("extract_pyclass_ref_mut", *span)
372-
} else {
373-
syn::Ident::new("extract_pyclass_ref", *span)
374-
};
375423
let holder = holders.push_holder(*span);
376424
let pyo3_path = pyo3_path.to_tokens_spanned(*span);
377-
error_mode.handle_error(
378-
quote_spanned! { *span =>
379-
#pyo3_path::impl_::extract_argument::#method::<#cls>(
380-
#arg,
381-
&mut #holder,
425+
match self_conversion.0 {
426+
SelfConversionPolicyInner::Trusted => {
427+
let method = if *mutable {
428+
syn::Ident::new("extract_pyclass_ref_mut_trusted", *span)
429+
} else {
430+
syn::Ident::new("extract_pyclass_ref_trusted", *span)
431+
};
432+
// Safety: slot wrappers are only installed on the extension type itself.
433+
// CPython's slot dispatch contract ensures the receiver is an instance
434+
// of the correct type before invoking the slot.
435+
//
436+
// The trailing `?` exists because if the extraction fails here it represents
437+
// a genuine type error, should not fall back to e.g. `ExtractErrorMode::NotImplemented`.
438+
quote! {
439+
unsafe { #pyo3_path::impl_::extract_argument::#method::<#cls>(
440+
#arg,
441+
&mut #holder,
442+
) }?
443+
}
444+
}
445+
SelfConversionPolicyInner::Checked => {
446+
let method = if *mutable {
447+
syn::Ident::new("extract_pyclass_ref_mut", *span)
448+
} else {
449+
syn::Ident::new("extract_pyclass_ref", *span)
450+
};
451+
error_mode.handle_error(
452+
quote_spanned! { *span =>
453+
#pyo3_path::impl_::extract_argument::#method::<#cls>(
454+
#arg,
455+
&mut #holder,
456+
)
457+
},
458+
ctx,
382459
)
383-
},
384-
ctx,
385-
)
460+
}
461+
}
386462
}
387463
SelfType::TryFromBoundRef { span, non_null } => {
388464
let bound_ref = if *non_null {
@@ -391,10 +467,34 @@ impl SelfType {
391467
quote! { unsafe { #pyo3_path::Bound::ref_from_ptr(#py, &#slf) } }
392468
};
393469
let pyo3_path = pyo3_path.to_tokens_spanned(*span);
470+
let receiver = match self_conversion.0 {
471+
SelfConversionPolicyInner::Trusted => {
472+
// Safety: slot wrappers are only installed on the extension type
473+
// itself. CPython's slot dispatch contract ensures the receiver is
474+
// an instance of the correct type (or a compatible subtype) before
475+
// invoking the slot.
476+
//
477+
// The wrapping `Ok(...?)` here is because an error here should not
478+
// be treated by e.g. `ExtractErrorMode::NotImplemented` as falling
479+
// back to the default, but instead a genuine type error.
480+
quote! {
481+
unsafe {
482+
#pyo3_path::PyResult::Ok(
483+
#pyo3_path::impl_::extract_argument::cast_bound_ref_trusted::<#cls>(#bound_ref)?
484+
)
485+
}
486+
}
487+
}
488+
SelfConversionPolicyInner::Checked => {
489+
quote_spanned! { *span =>
490+
#bound_ref.cast::<#cls>()
491+
.map_err(::std::convert::Into::<#pyo3_path::PyErr>::into)
492+
}
493+
}
494+
};
394495
error_mode.handle_error(
395496
quote_spanned! { *span =>
396-
#bound_ref.cast::<#cls>()
397-
.map_err(::std::convert::Into::<#pyo3_path::PyErr>::into)
497+
#receiver
398498
.and_then(
399499
#[allow(
400500
clippy::unnecessary_fallible_conversions,
@@ -678,6 +778,7 @@ impl<'a> FnSpec<'a> {
678778
ident: &proc_macro2::Ident,
679779
cls: Option<&syn::Type>,
680780
convention: CallingConvention,
781+
self_conversion: SelfConversionPolicy,
681782
ctx: &Ctx,
682783
) -> Result<TokenStream> {
683784
let Ctx {
@@ -700,9 +801,13 @@ impl<'a> FnSpec<'a> {
700801
}
701802

702803
let rust_call = |args: Vec<TokenStream>, mut holders: Holders| {
703-
let self_arg = self
704-
.tp
705-
.self_arg(cls, ExtractErrorMode::Raise, &mut holders, ctx);
804+
let self_arg = self.tp.self_arg(
805+
cls,
806+
ExtractErrorMode::Raise,
807+
self_conversion,
808+
&mut holders,
809+
ctx,
810+
);
706811
let init_holders = holders.init_holders(ctx);
707812

708813
// We must assign the output_span to the return value of the call,

‎pyo3-macros-backend/src/pyfunction.rs‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ use crate::{
1010
self, get_pyo3_options, take_attributes, take_pyo3_options, CrateAttribute,
1111
FromPyWithAttribute, NameAttribute, TextSignatureAttribute,
1212
},
13-
method::{self, CallingConvention, FnArg},
13+
method::{self, CallingConvention, FnArg, SelfConversionPolicy},
1414
pymethod::check_generic,
1515
};
1616
use proc_macro2::{Span, TokenStream};
@@ -430,7 +430,13 @@ pub fn impl_wrap_pyfunction(
430430
);
431431
}
432432
let calling_convention = CallingConvention::from_signature(&spec.signature);
433-
let wrapper = spec.get_wrapper_function(&wrapper_ident, None, calling_convention, ctx)?;
433+
let wrapper = spec.get_wrapper_function(
434+
&wrapper_ident,
435+
None,
436+
calling_convention,
437+
SelfConversionPolicy::checked(),
438+
ctx,
439+
)?;
434440
let methoddef = spec.get_methoddef(
435441
wrapper_ident,
436442
spec.get_doc(&func.attrs).as_ref(),

0 commit comments

Comments
 (0)