From b582f3b9bb544339ef2719a3651873c0aba9ef3e Mon Sep 17 00:00:00 2001
From: Kyle Barron <kylebarron2@gmail.com>
Date: Wed, 31 Jul 2024 12:17:54 -0400
Subject: [PATCH] Make readers python iterators (#85)

---
 arro3-core/python/arro3/core/_core.pyi |  4 ++++
 pyo3-arrow/src/array_reader.rs         | 10 ++++++++++
 pyo3-arrow/src/record_batch_reader.rs  | 10 ++++++++++
 3 files changed, 24 insertions(+)

diff --git a/arro3-core/python/arro3/core/_core.pyi b/arro3-core/python/arro3/core/_core.pyi
index 37a6337..8856d94 100644
--- a/arro3-core/python/arro3/core/_core.pyi
+++ b/arro3-core/python/arro3/core/_core.pyi
@@ -73,6 +73,8 @@ class Array:
 
 class ArrayReader:
     def __arrow_c_stream__(self, requested_schema: object | None = None) -> object: ...
+    def __iter__(self) -> ArrayReader: ...
+    def __next__(self) -> Array: ...
     def __repr__(self) -> str: ...
     @classmethod
     def from_arrow(
@@ -725,6 +727,8 @@ class RecordBatch:
 
 class RecordBatchReader:
     def __arrow_c_stream__(self, requested_schema: object | None = None) -> object: ...
+    def __iter__(self) -> RecordBatchReader: ...
+    def __next__(self) -> RecordBatch: ...
     def __repr__(self) -> str: ...
     @classmethod
     def from_arrow(
diff --git a/pyo3-arrow/src/array_reader.rs b/pyo3-arrow/src/array_reader.rs
index 85ff2e1..e4aa202 100644
--- a/pyo3-arrow/src/array_reader.rs
+++ b/pyo3-arrow/src/array_reader.rs
@@ -122,6 +122,16 @@ impl PyArrayReader {
         to_stream_pycapsule(py, array_reader, requested_schema)
     }
 
+    // Return self
+    // https://stackoverflow.com/a/52056290
+    fn __iter__(&mut self, py: Python) -> PyResult<PyObject> {
+        self.to_arro3(py)
+    }
+
+    fn __next__(&mut self, py: Python) -> PyArrowResult<PyObject> {
+        self.read_next_array(py)
+    }
+
     pub fn __repr__(&self) -> String {
         self.to_string()
     }
diff --git a/pyo3-arrow/src/record_batch_reader.rs b/pyo3-arrow/src/record_batch_reader.rs
index 8719220..66bd3ab 100644
--- a/pyo3-arrow/src/record_batch_reader.rs
+++ b/pyo3-arrow/src/record_batch_reader.rs
@@ -145,6 +145,16 @@ impl PyRecordBatchReader {
         to_stream_pycapsule(py, array_reader, requested_schema)
     }
 
+    // Return self
+    // https://stackoverflow.com/a/52056290
+    fn __iter__(&mut self, py: Python) -> PyResult<PyObject> {
+        self.to_arro3(py)
+    }
+
+    fn __next__(&mut self, py: Python) -> PyArrowResult<PyObject> {
+        self.read_next_batch(py)
+    }
+
     pub fn __repr__(&self) -> String {
         self.to_string()
     }