diff --git a/python/src/cluster_file.hpp b/python/src/cluster_file.hpp index 0b6fbb06..1801924b 100644 --- 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_(m, "ClusterFile") - .def(py::init()) + .def(py::init(), py::arg(), py::arg("chunk_size") = 1000) .def("read_clusters", [](ClusterFile &self, size_t n_clusters) { auto* vec = new std::vector(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(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(self.read_clusters(self.chunk_size())); + if(vec->size() == 0) { + throw py::stop_iteration(); + } + return return_vector(vec); + }); } \ No newline at end of file diff --git a/src/ClusterFile.cpp b/src/ClusterFile.cpp index f0aa1736..aff4625a 100644 --- 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());