Add docstrings and convert to properties

anjakefala · Jul 24, 2024 · 9c10731 · 9c10731
1 parent d5b4b9d
commit 9c10731
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 3 deletions.
diff --git a/python/pyarrow/_cuda.pyx b/python/pyarrow/_cuda.pyx
@@ -185,10 +185,27 @@ cdef class Context(_Weakrefable):
             cudabuf = GetResultValue(self.context.get().Allocate(nbytes))
         return pyarrow_wrap_cudabuffer(cudabuf)
 
+    @property
     def memory_manager(self):
+        """
+        The default memory manager tied to this context's device.
+
+        Returns
+        -------
+        MemoryManager
+        """
         return MemoryManager.wrap(self.context.get().memory_manager())
 
+    @property
     def device(self):
+        """
+        The device instance associated with this context.
+
+
+        Returns
+        -------
+        Device
+        """
         return Device.wrap(self.context.get().device())
 
     def foreign_buffer(self, address, size, base=None):

diff --git a/python/pyarrow/includes/libarrow_cuda.pxd b/python/pyarrow/includes/libarrow_cuda.pxd
@@ -41,8 +41,8 @@ cdef extern from "arrow/gpu/cuda_api.h" namespace "arrow::cuda" nogil:
         const void* handle() const
         int device_number() const
         CResult[uintptr_t] GetDeviceAddress(uintptr_t addr)
-        shared_ptr[CDevice] device() const
-        shared_ptr[CMemoryManager] memory_manager() const
+        shared_ptr[CDevice] device const
+        shared_ptr[CMemoryManager] memory_manager const
 
     cdef cppclass CCudaIpcMemHandle" arrow::cuda::CudaIpcMemHandle":
         @staticmethod

diff --git a/python/pyarrow/tests/test_cuda.py b/python/pyarrow/tests/test_cuda.py
@@ -67,7 +67,6 @@ def test_Context():
     assert not dev.is_cpu
     assert dev.device_id == 0
     assert dev.device_type == pa.DeviceAllocationType.CUDA
-    assert "<pyarrow.Device: CudaDevice" in repr(dev)
 
     with pytest.raises(ValueError,
                        match=("device_number argument must "