Make buffers have a length (#252)

* Make buffers have a length * Update pyo3-arrow/src/buffer.rs Co-authored-by: Kyle Barron <[email protected]> * Add `__len__` to typing --------- Co-authored-by: Kyle Barron <[email protected]> Co-authored-by: Kyle Barron <[email protected]>
kylebarron · Oct 30, 2024 · 2c73b76 · 2c73b76
1 parent d771e39
commit 2c73b76
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 0 deletions.
diff --git a/arro3-core/python/arro3/core/_core.pyi b/arro3-core/python/arro3/core/_core.pyi
@@ -210,6 +210,7 @@ class Buffer:
     """An Arrow Buffer"""
     def __init__(self, buffer) -> None: ...
     def __buffer__(self, flags: int) -> memoryview: ...
+    def __len__(self) -> int: ...
     def as_bytes(self) -> bytes:
         """Copy this buffer into a Python `bytes` object."""
 

diff --git a/pyo3-arrow/src/buffer.rs b/pyo3-arrow/src/buffer.rs
@@ -80,6 +80,10 @@ impl PyArrowBuffer {
         PyBytes::new_bound(py, &self.0)
     }
 
+    fn __len__(&self) -> usize {
+        self.0.len()
+    }
+
     /// This is taken from opendal:
     /// https://github.com/apache/opendal/blob/d001321b0f9834bc1e2e7d463bcfdc3683e968c9/bindings/python/src/utils.rs#L51-L72
     unsafe fn __getbuffer__(

diff --git a/tests/core/test_buffer_protocol.py b/tests/core/test_buffer_protocol.py
@@ -71,6 +71,7 @@ def test_multi_dimensional():
 def test_round_trip_buffer():
     arr = np.arange(5, dtype=np.uint8)
     buffer = Buffer(arr)
+    assert len(buffer) == arr.nbytes
     retour = np.frombuffer(buffer, dtype=np.uint8)
     assert np.array_equal(arr, retour)