Skip to content

Commit

Permalink
chore: add a warning if failed to get allocator capsule
Browse files Browse the repository at this point in the history
  • Loading branch information
ruihe774 committed Jul 26, 2024
1 parent d17b225 commit 98ef808
Showing 1 changed file with 9 additions and 4 deletions.
13 changes: 9 additions & 4 deletions pyo3-polars/src/alloc.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
use std::alloc::{GlobalAlloc, Layout, System};
use std::ffi::c_char;

use once_cell::race::OnceRef;
use pyo3::ffi::{PyCapsule_Import, Py_IsInitialized};
use pyo3::Python;
Expand Down Expand Up @@ -67,16 +66,22 @@ impl PolarsAllocator {
// Do not allocate in this function,
// otherwise it will cause infinite recursion.
self.0.get_or_init(|| {
(unsafe { Py_IsInitialized() } != 0)
let r = (unsafe { Py_IsInitialized() } != 0)
.then(|| {
Python::with_gil(|_| unsafe {
(PyCapsule_Import(ALLOCATOR_CAPSULE_NAME.as_ptr() as *const c_char, 0)
as *const AllocatorCapsule)
.as_ref()
})
})
.flatten()
.unwrap_or(&FALLBACK_ALLOCATOR_CAPSULE)
.flatten();
#[cfg(debug_assertions)]
if r.is_none() {
// Do not use eprintln; it may alloc
use std::io::Write;
std::io::stderr().write_all(b"failed to get allocator capsule\n").unwrap();
}
r.unwrap_or(&FALLBACK_ALLOCATOR_CAPSULE)
})
}

Expand Down

0 comments on commit 98ef808

Please sign in to comment.