WIP

slsdetectorgroup · Nov 14, 2024 · 0d05827 · 0d05827
1 parent 5cde7a9
commit 0d05827
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/python/src/cluster_file.hpp b/python/src/cluster_file.hpp
@@ -18,7 +18,7 @@ void define_cluster_file_io_bindings(py::module &m) {
     PYBIND11_NUMPY_DTYPE(Cluster, x, y, data);
 
     py::class_<ClusterFile>(m, "ClusterFile")
-        .def(py::init<const std::filesystem::path &>())
+        .def(py::init<const std::filesystem::path &, size_t>(), py::arg(), py::arg("chunk_size") = 1000)
         .def("read_clusters",
              [](ClusterFile &self, size_t n_clusters) {
                  auto* vec = new std::vector<Cluster>(self.read_clusters(n_clusters));
@@ -29,6 +29,16 @@ void define_cluster_file_io_bindings(py::module &m) {
                  auto view = make_view_2d(noise_map);
                  auto* vec = new std::vector<Cluster>(self.read_cluster_with_cut(n_clusters, view.data(), nx, ny));
                  return return_vector(vec);
-             });
+             })
+        .def("__enter__", [](ClusterFile &self) { return &self; })
+        .def("__exit__", [](ClusterFile &self, py::args args) { return; })
+        .def("__iter__", [](ClusterFile &self) { return &self; })
+        .def("__next__", [](ClusterFile &self) {
+            auto vec =  new std::vector<Cluster>(self.read_clusters(self.chunk_size()));
+            if(vec->size() == 0) {
+                throw py::stop_iteration();
+            }
+            return return_vector(vec);
+        });
 
 }
diff --git a/src/ClusterFile.cpp b/src/ClusterFile.cpp
@@ -2,7 +2,7 @@
 
 namespace aare {
 
-ClusterFile::ClusterFile(const std::filesystem::path &fname) {
+ClusterFile::ClusterFile(const std::filesystem::path &fname, size_t chunk_size): m_chunk_size(chunk_size) {
     fp = fopen(fname.c_str(), "rb");
     if (!fp) {
         throw std::runtime_error("Could not open file: " + fname.string());