diff --git a/src/types/sheet.rs b/src/types/sheet.rs index 5042a6d..5c4c0af 100644 --- 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, ) -> PyResult>> { - 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 index 5aa5024..d8e6b25 100644 --- 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"],