Merge pull request #37 from dimastbk/fix-nrows

fix nrows when skip_empty_area=False; increase performance for some c…
dimastbk · Sep 12, 2023 · a690b3b · a690b3b
2 parents b2ee2cf + 2229c61
commit a690b3b
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 20 deletions.
diff --git a/src/types/sheet.rs b/src/types/sheet.rs
@@ -1,3 +1,4 @@
+use std::borrow::Cow;
 use std::fmt::Display;
 
 use calamine::{DataType, Range, SheetType, SheetVisible};
@@ -168,28 +169,25 @@ impl CalamineSheet {
         skip_empty_area: bool,
         nrows: Option<u32>,
     ) -> PyResult<Vec<Vec<CellValue>>> {
-        let mut range = self.range.to_owned();
-
-        if !skip_empty_area {
-            if let Some(end) = range.end() {
-                range = range.range((0, 0), end)
-            }
-        }
-
-        if let Some(nrows) = nrows {
-            if range.end().is_some() && range.start().is_some() {
-                range = range.range(
-                    range.start().unwrap(),
-                    (
-                        range.start().unwrap().0 + (nrows - 1),
-                        range.end().unwrap().1,
-                    ),
-                )
-            }
-        }
+        let nrows = match nrows {
+            Some(nrows) => nrows,
+            None => self.range.end().map_or(0, |end| end.0 + 1),
+        };
+
+        let range = if skip_empty_area {
+            Cow::Borrowed(&self.range)
+        } else if let Some(end) = self.range.end() {
+            Cow::Owned(self.range.range(
+                (0, 0),
+                (if nrows > end.0 { end.0 } else { nrows - 1 }, end.1),
+            ))
+        } else {
+            Cow::Borrowed(&self.range)
+        };
 
         Ok(range
             .rows()
+            .take(nrows as usize)
             .map(|row| row.iter().map(|x| x.into()).collect())
             .collect())
     }

diff --git a/tests/test_base.py b/tests/test_base.py
@@ -184,18 +184,33 @@ def test_nrows():
     reader = CalamineWorkbook.from_object(PATH / "base.xlsx")
     sheet = reader.get_sheet_by_name("Sheet3")
 
-    assert sheet.to_python(nrows=1) == [["line1", "line1", "line1"]]
+    assert sheet.to_python(nrows=1) == [
+        ["line1", "line1", "line1"],
+    ]
 
     assert sheet.to_python(nrows=2) == [
         ["line1", "line1", "line1"],
         ["line2", "line2", "line2"],
     ]
 
+    assert sheet.to_python(nrows=4) == [
+        ["line1", "line1", "line1"],
+        ["line2", "line2", "line2"],
+        ["line3", "line3", "line3"],
+    ]
+
     assert sheet.to_python(skip_empty_area=False, nrows=2) == [
         ["", "", "", ""],
         ["", "line1", "line1", "line1"],
     ]
 
+    assert sheet.to_python(skip_empty_area=False, nrows=5) == [
+        ["", "", "", ""],
+        ["", "line1", "line1", "line1"],
+        ["", "line2", "line2", "line2"],
+        ["", "line3", "line3", "line3"],
+    ]
+
     assert sheet.to_python() == [
         ["line1", "line1", "line1"],
         ["line2", "line2", "line2"],