Add datareader.size

ecmwf · Nov 18, 2024 · 69c5d9a · 69c5d9a
1 parent 0d6c2a4
commit 69c5d9a
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 4 deletions.
diff --git a/src/pyfdb/processed_fdb.h b/src/pyfdb/processed_fdb.h
@@ -45,6 +45,7 @@ int fdb_datareader_close(fdb_datareader_t* dr);
 int fdb_datareader_tell(fdb_datareader_t* dr, long* pos);
 int fdb_datareader_seek(fdb_datareader_t* dr, long pos);
 int fdb_datareader_skip(fdb_datareader_t* dr, long count);
+int fdb_datareader_size(fdb_datareader_t* dr, long* size);
 int fdb_datareader_read(fdb_datareader_t* dr, void *buf, long count, long* read);
 int fdb_delete_datareader(fdb_datareader_t* dr);
 

diff --git a/src/pyfdb/pyfdb.py b/src/pyfdb/pyfdb.py
@@ -267,12 +267,20 @@ def tell(self):
         lib.fdb_datareader_tell(self.__dataread, where)
         return where[0]
 
-    def read(self, count=-1) -> bytes:
+    def size(self):
+        size = ffi.new("long*")
+        lib.fdb_datareader_size(self.__dataread, size)
+        return size[0]
+
+
+    def read(self, size=-1) -> bytes:
         self.open()
-        if isinstance(count, int):
-            buf = bytearray(count)
+        if isinstance(size, int):
+            if (size == -1 ): 
+                size = self.size()
+            buf = bytearray(size)
             read = ffi.new("long*")
-            lib.fdb_datareader_read(self.__dataread, ffi.from_buffer(buf), count, read)
+            lib.fdb_datareader_read(self.__dataread, ffi.from_buffer(buf), size, read)
             return buf[0 : read[0]]
         return bytearray()