fix: Update MemoryReader to read more than one field

eccodes/highlevel/reader.py

@@ -74,12 +74,21 @@ class MemoryReader(ReaderBase):
     def __init__(self, buf, kind=eccodes.CODES_PRODUCT_GRIB):
         super().__init__(kind=kind)
         self.buf = buf
+        self._index = 0
 
     def _next_handle(self):
         if self.buf is None:
             return None
-        handle = eccodes.codes_new_from_message(self.buf)
-        self.buf = None
+        handle = eccodes.codes_new_from_message(self.buf[self._index :])
+
+        try:
+            handle_length = eccodes.codes_get(handle, "totalLength")
+        except eccodes.GribInternalError:
+            return None
+
+        self._index += handle_length
+        if self._index >= len(self.buf):
+            self.buf = None
         return handle
 
 

tests/test_highlevel.py

@@ -1,4 +1,5 @@
 import collections
+import io
 import itertools
 import pathlib
 
@@ -179,3 +180,63 @@ def test_grib_message_from_samples():
 def test_bufr_message_from_samples():
     message = eccodes.BUFRMessage.from_samples("BUFR4")
     assert message["edition"] == 4
+
+
+def test_read_memory():
+    buffer = io.BytesIO()
+    with eccodes.FileReader(TEST_GRIB_DATA2) as reader1:
+        written = []
+        for message in itertools.islice(reader1, 15):
+            message.write_to(buffer)
+            written.append(message)
+
+    with eccodes.MemoryReader(buffer.getvalue()) as reader2:
+        for message1, message2 in itertools.zip_longest(written, reader2):
+            assert message1 is not None
+            assert message2 is not None
+            for key in [
+                "edition",
+                "centre",
+                "typeOfLevel",
+                "level",
+                "dataDate",
+                "stepRange",
+                "dataType",
+                "shortName",
+                "packingType",
+                "gridType",
+                "number",
+            ]:
+                assert message1[key] == message2[key]
+            assert np.all(message1.data == message2.data)
+
+
+def test_read_stream():
+    buffer = io.BytesIO()
+    with eccodes.FileReader(TEST_GRIB_DATA2) as reader1:
+        written = []
+        for message in itertools.islice(reader1, 15):
+            message.write_to(buffer)
+            written.append(message)
+
+    buffer.seek(0, 0)
+
+    with eccodes.StreamReader(buffer) as reader2:
+        for message1, message2 in itertools.zip_longest(written, reader2):
+            assert message1 is not None
+            assert message2 is not None
+            for key in [
+                "edition",
+                "centre",
+                "typeOfLevel",
+                "level",
+                "dataDate",
+                "stepRange",
+                "dataType",
+                "shortName",
+                "packingType",
+                "gridType",
+                "number",
+            ]:
+                assert message1[key] == message2[key]
+            assert np.all(message1.data == message2.data)