Merge pull request #43 from dimastbk/issue42

feat: iterator over rust range
dimastbk · Jul 15, 2024 · 8765389 · 8765389
2 parents 83523e5 + ba8e6dc
commit 8765389
Show file tree

Hide file tree

Showing 4 changed files with 94 additions and 2 deletions.
diff --git a/python/python_calamine/_python_calamine.pyi b/python/python_calamine/_python_calamine.pyi
@@ -72,6 +72,22 @@ class CalamineSheet:
             For suppress this behaviour, set `skip_empty_area` to `False`.
         """
 
+    def iter_rows(
+        self,
+    ) -> typing.Iterator[
+        list[
+            int
+            | float
+            | str
+            | bool
+            | datetime.time
+            | datetime.date
+            | datetime.datetime
+            | datetime.timedelta
+        ]
+    ]:
+        """Retunrning data from sheet as iterator of lists."""
+
 @typing.final
 class CalamineWorkbook:
     path: str | None

diff --git a/src/types/cell.rs b/src/types/cell.rs
@@ -3,7 +3,7 @@ use std::convert::From;
 use calamine::DataType;
 use pyo3::prelude::*;
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub enum CellValue {
     Int(i64),
     Float(f64),

diff --git a/src/types/sheet.rs b/src/types/sheet.rs
@@ -1,7 +1,7 @@
 use std::fmt::Display;
 use std::sync::Arc;
 
-use calamine::{Data, Range, SheetType, SheetVisible};
+use calamine::{Data, Range, Rows, SheetType, SheetVisible};
 use pyo3::class::basic::CompareOp;
 use pyo3::prelude::*;
 use pyo3::types::PyList;
@@ -196,4 +196,56 @@ impl CalamineSheet {
             }),
         ))
     }
+
+    fn iter_rows(&self) -> CalamineCellIterator {
+        CalamineCellIterator::from_range(Arc::clone(&self.range))
+    }
+}
+
+#[pyclass]
+pub struct CalamineCellIterator {
+    position: u32,
+    start: (u32, u32),
+    empty_row: Vec<CellValue>,
+    iter: Rows<'static, Data>,
+    #[allow(dead_code)]
+    range: Arc<Range<Data>>,
+}
+
+impl CalamineCellIterator {
+    fn from_range(range: Arc<Range<Data>>) -> CalamineCellIterator {
+        let empty_row = (0..range.width())
+            .map(|_| CellValue::String("".to_string()))
+            .collect();
+        CalamineCellIterator {
+            empty_row,
+            position: 0,
+            start: range.start().unwrap(),
+            iter: unsafe {
+                std::mem::transmute::<
+                    calamine::Rows<'_, calamine::Data>,
+                    calamine::Rows<'static, calamine::Data>,
+                >(range.rows())
+            },
+            range,
+        }
+    }
+}
+
+#[pymethods]
+impl CalamineCellIterator {
+    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
+        slf
+    }
+
+    fn __next__(mut slf: PyRefMut<'_, Self>) -> Option<Bound<'_, PyList>> {
+        slf.position += 1;
+        if slf.position > slf.start.0 {
+            slf.iter.next().map(|row| {
+                PyList::new_bound(slf.py(), row.iter().map(<&Data as Into<CellValue>>::into))
+            })
+        } else {
+            Some(PyList::new_bound(slf.py(), slf.empty_row.clone()))
+        }
+    }
 }
diff --git a/tests/test_base.py b/tests/test_base.py
@@ -180,6 +180,30 @@ def test_xlsx_read():
     assert [] == reader.get_sheet_by_index(1).to_python(skip_empty_area=False)
 
 
+def test_xlsx_iter_rows():
+    names = ["Sheet1", "Sheet2", "Sheet3"]
+    data = [
+        ["", "", "", "", "", "", "", "", "", ""],
+        [
+            "String",
+            1,
+            1.1,
+            True,
+            False,
+            date(2010, 10, 10),
+            datetime(2010, 10, 10, 10, 10, 10),
+            time(10, 10, 10),
+            timedelta(hours=10, minutes=10, seconds=10, microseconds=100000),
+            timedelta(hours=255, minutes=10, seconds=10),
+        ],
+    ]
+
+    reader = CalamineWorkbook.from_object(PATH / "base.xlsx")
+
+    assert names == reader.sheet_names
+    assert data == list(reader.get_sheet_by_index(0).iter_rows())
+
+
 def test_nrows():
     reader = CalamineWorkbook.from_object(PATH / "base.xlsx")
     sheet = reader.get_sheet_by_name("Sheet3")