From 98ef808853bdc105efa9e40bfbd70b5f562bc759 Mon Sep 17 00:00:00 2001 From: Misaki Kasumi Date: Fri, 26 Jul 2024 10:31:01 +0800 Subject: [PATCH] chore: add a warning if failed to get allocator capsule --- pyo3-polars/src/alloc.rs | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/pyo3-polars/src/alloc.rs b/pyo3-polars/src/alloc.rs index 97ea645..a429704 100644 --- a/pyo3-polars/src/alloc.rs +++ b/pyo3-polars/src/alloc.rs @@ -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; @@ -67,7 +66,7 @@ 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) @@ -75,8 +74,14 @@ impl PolarsAllocator { .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) }) }