From ba909d6f16293e944fa0134700ab110b94425a4f Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Fri, 25 Oct 2024 11:15:28 -0700 Subject: [PATCH] address review --- python/pylibcudf/pylibcudf/interop.pyx | 14 +++++++++----- python/pylibcudf/pylibcudf/tests/test_interop.py | 2 +- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/python/pylibcudf/pylibcudf/interop.pyx b/python/pylibcudf/pylibcudf/interop.pyx index ba50010bc3a..61e812353b7 100644 --- a/python/pylibcudf/pylibcudf/interop.pyx +++ b/python/pylibcudf/pylibcudf/interop.pyx @@ -342,16 +342,18 @@ cpdef Table from_dlpack(object managed_tensor): Table with a copy of the tensor data. """ if not PyCapsule_IsValid(managed_tensor, "dltensor"): - raise ValueError("Invalid capsule object") + raise ValueError("Invalid PyCapsule object") cdef unique_ptr[table] c_result cdef DLManagedTensor* dlpack_tensor = PyCapsule_GetPointer( managed_tensor, "dltensor" ) + if dlpack_tensor is NULL: + raise ValueError("PyCapsule object contained a NULL pointer") PyCapsule_SetName(managed_tensor, "used_dltensor") # Note: A copy is always performed when converting the dlpack # data to a libcudf table. We also delete the dlpack_tensor pointer - # as the poionter is not deleted by libcudf's from_dlpack function. + # as the pointer is not deleted by libcudf's from_dlpack function. # TODO: https://github.com/rapidsai/cudf/issues/10874 # TODO: https://github.com/rapidsai/cudf/issues/10849 with nogil: @@ -400,6 +402,8 @@ cdef void dlmanaged_tensor_pycapsule_deleter(object pycap_obj) noexcept: if PyCapsule_IsValid(pycap_obj, "used_dltensor"): # we do not call a used capsule's deleter return - cdef DLManagedTensor* dlpack_tensor - dlpack_tensor = PyCapsule_GetPointer(pycap_obj, "dltensor") - dlpack_tensor.deleter(dlpack_tensor) + cdef DLManagedTensor* dlpack_tensor = PyCapsule_GetPointer( + pycap_obj, "dltensor" + ) + if dlpack_tensor is not NULL: + dlpack_tensor.deleter(dlpack_tensor) diff --git a/python/pylibcudf/pylibcudf/tests/test_interop.py b/python/pylibcudf/pylibcudf/tests/test_interop.py index 41d1b24f740..654f0ee4d32 100644 --- a/python/pylibcudf/pylibcudf/tests/test_interop.py +++ b/python/pylibcudf/pylibcudf/tests/test_interop.py @@ -95,5 +95,5 @@ def test_to_dlpack_error(): def test_from_dlpack_error(): - with pytest.raises(ValueError, match="Invalid capsule object"): + with pytest.raises(ValueError, match="Invalid PyCapsule object"): plc.interop.from_dlpack(1)