make iterator repeatable

pythonApi/grib_to_arrow.cpp

@@ -109,6 +109,9 @@ PYBIND11_MODULE(gribtoarrow, m)
             e.g. If you wanted to convert from Kelvin to Celcius you would pass a table which contained the parameterId and
             contained 273.15 in the column subtraction_value                
         )EOL") 
+        .def("withRepeatableIterator", &GribReader::withRepeatableIterator, pybind11::call_guard<pybind11::gil_scoped_release>(), R"EOL(
+            Enables the message to be iterated multiple times.                 
+        )EOL") 
         .def(
             "__iter__",
             [](GribReader &s) { return py::make_iterator(s.begin(), s.end()); },

src/gribreader.cpp

@@ -253,6 +253,7 @@ optional<function<arrow::Result<shared_ptr<arrow::Array>>(shared_ptr<arrow::Arra
 }
 
 Iterator GribReader::begin() { 
+    std::cout << "Creating iterator" << endl;
     codes_handle* h = codes_handle_new_from_file(0, fin, PRODUCT_GRIB, &err);
     auto m = new GribMessage(this, h, 0l);
     return Iterator(this, m, m_endMessage);
@@ -261,8 +262,8 @@ Iterator GribReader::begin() {
 Iterator GribReader::end()   {
 
     if (isRepeatable) {
-        //TODO Make the iterator repeatable
-        //Not sure how yet need to play with how this iterfaces at the python __iter__, next() and StopIteration level
+        std::cout << "At end of message and is repeatable" << endl;
+        fseek(fin, 0, SEEK_SET);
     }
     return Iterator( this,  m_endMessage, m_endMessage );
 

tests/test_iteration.py

@@ -0,0 +1,33 @@
+import polars as pl
+import pytest
+import os
+import sys
+import pyarrow
+from functools import reduce
+
+class TestIteration:
+
+    def get_iterator_count(self, reader):
+        return sum(1 for _ in reader)
+
+    def test_iterate_once(self, resource):
+        from gribtoarrow import GribReader
+
+        reader = GribReader(str(resource) + "/meps_weatherapi_sorlandet.grb")
+
+        cnt = self.get_iterator_count(reader)
+        assert cnt == 269
+        cnt = self.get_iterator_count(reader)
+        assert cnt == 0
+
+    def test_iterate_once(self, resource):
+        from gribtoarrow import GribReader
+
+        reader = GribReader(str(resource) + "/meps_weatherapi_sorlandet.grb").withRepeatableIterator(True)
+
+        it = iter(reader)
+
+        cnt = self.get_iterator_count(reader)
+        assert cnt == 269
+        cnt = self.get_iterator_count(reader)
+        assert cnt == 269