From d17b225e2c30d81c7ccabe9fb4242130edfb5975 Mon Sep 17 00:00:00 2001 From: Misaki Kasumi Date: Fri, 26 Jul 2024 09:10:39 +0800 Subject: [PATCH] fix: check Py_IsInitialized before acquire GIL in PolarsAllocator --- pyo3-polars/src/alloc.rs | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/pyo3-polars/src/alloc.rs b/pyo3-polars/src/alloc.rs index 499b3f9..97ea645 100644 --- a/pyo3-polars/src/alloc.rs +++ b/pyo3-polars/src/alloc.rs @@ -2,7 +2,7 @@ use std::alloc::{GlobalAlloc, Layout, System}; use std::ffi::c_char; use once_cell::race::OnceRef; -use pyo3::ffi::PyCapsule_Import; +use pyo3::ffi::{PyCapsule_Import, Py_IsInitialized}; use pyo3::Python; unsafe extern "C" fn fallback_alloc(size: usize, align: usize) -> *mut u8 { @@ -67,14 +67,16 @@ impl PolarsAllocator { // Do not allocate in this function, // otherwise it will cause infinite recursion. self.0.get_or_init(|| { - Python::with_gil(|_| { - unsafe { - (PyCapsule_Import(ALLOCATOR_CAPSULE_NAME.as_ptr() as *const c_char, 0) - as *const AllocatorCapsule) - .as_ref() - } + (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) - }) }) }