cpp_RemoteEndpoint

rapidsai · Sep 20, 2024 · 6398ddc · 6398ddc
1 parent 783ec08
commit 6398ddc
Showing 1 changed file with 25 additions and 11 deletions.
diff --git a/python/kvikio/kvikio/_lib/remote_handle.pyx b/python/kvikio/kvikio/_lib/remote_handle.pyx
@@ -10,19 +10,24 @@ from cython.operator cimport dereference as deref
 from libc.stdint cimport uintptr_t
 from libcpp.memory cimport make_unique, unique_ptr
 from libcpp.string cimport string
-from libcpp.utility cimport pair
+from libcpp.utility cimport move, pair
 
 from kvikio._lib.arr cimport parse_buffer_argument
 from kvikio._lib.future cimport IOFuture, _wrap_io_future, future
 
 
 cdef extern from "<kvikio/remote_handle.hpp>" nogil:
+    cdef cppclass cpp_RemoteEndpoint "kvikio::RemoteEndpoint":
+        pass
+
     cdef cppclass cpp_HttpEndpoint "kvikio::HttpEndpoint":
         cpp_HttpEndpoint(string url) except +
 
     cdef cppclass cpp_RemoteHandle "kvikio::RemoteHandle":
-        cpp_RemoteHandle(cpp_HttpEndpoint endpoint, size_t nbytes) except +
-        cpp_RemoteHandle(cpp_HttpEndpoint endpoint) except +
+        cpp_RemoteHandle(
+            unique_ptr[cpp_RemoteEndpoint] endpoint, size_t nbytes
+        ) except +
+        cpp_RemoteHandle(unique_ptr[cpp_RemoteEndpoint] endpoint) except +
         int nbytes() except +
         size_t read(
             void* buf,
@@ -35,22 +40,31 @@ cdef extern from "<kvikio/remote_handle.hpp>" nogil:
             size_t file_offset
         ) except +
 
+cdef string _to_string(str_or_none):
+    """Convert Python object to a C++ string (if None, return the empty string)"""
+    if str_or_none is None:
+        return string()
+    return str.encode(str(str_or_none))
+
+
 cdef class RemoteFile:
     cdef unique_ptr[cpp_RemoteHandle] _handle
 
     @classmethod
-    def from_url(cls, url: str, nbytes: Optional[int]):
+    def from_url(
+        cls,
+        url: str,
+        nbytes: Optional[int],
+    ):
         cdef RemoteFile ret = RemoteFile()
-        cdef string u = str.encode(str(url))
+        cdef unique_ptr[cpp_HttpEndpoint] ep = make_unique[cpp_HttpEndpoint](
+            _to_string(url)
+        )
         if nbytes is None:
-            ret._handle = make_unique[cpp_RemoteHandle](
-                make_unique[cpp_HttpEndpoint](u)
-            )
+            ret._handle = make_unique[cpp_RemoteHandle](move(ep))
             return ret
         cdef size_t n = nbytes
-        ret._handle = make_unique[cpp_RemoteHandle](
-            make_unique[cpp_HttpEndpoint](u), n
-        )
+        ret._handle = make_unique[cpp_RemoteHandle](move(ep), n)
         return ret
 
     def nbytes(self) -> int: