Repository: dimastbk/python-calamine (language: python)

Issue:
No way to close a CalamineWorkbook?
There is no .close() method and as far as I can tell a with block has no effect either. Can these features please be implemented so the file being read doesn't remain locked until the executing program closes?


Discussion / hints:
Hi. Could you give a reproducible example?
A reproducible example of a missing feature?

Sure. Create a CalamineWorkbook then try to close it.

Fix patch:
diff --git a/python/python_calamine/__init__.py b/python/python_calamine/__init__.py
index 9771394..d2f57ea 100644
--- a/python/python_calamine/__init__.py
+++ b/python/python_calamine/__init__.py
@@ -6,6 +6,7 @@
     SheetMetadata,
     SheetTypeEnum,
     SheetVisibleEnum,
+    WorkbookClosed,
     WorksheetNotFound,
     XmlError,
     ZipError,
@@ -23,5 +24,6 @@
     "WorksheetNotFound",
     "XmlError",
     "ZipError",
+    "WorkbookClosed",
     "load_workbook",
 )
diff --git a/python/python_calamine/_python_calamine.pyi b/python/python_calamine/_python_calamine.pyi
index cb9688c..ad92182 100644
--- a/python/python_calamine/_python_calamine.pyi
+++ b/python/python_calamine/_python_calamine.pyi
@@ -1,8 +1,10 @@
 from __future__ import annotations
 
+import contextlib
 import datetime
 import enum
 import os
+import types
 import typing
 
 @typing.type_check_only
@@ -65,11 +67,10 @@ class CalamineSheet:
     ]:
         """Retunrning data from sheet as list of lists.
 
-        Parameters
-        ----------
-        skip_empty_area : bool
-            By default, calamine skips empty rows/cols before data.
-            For suppress this behaviour, set `skip_empty_area` to `False`.
+        Args:
+            skip_empty_area (bool):
+                By default, calamine skips empty rows/cols before data.
+                For suppress this behaviour, set `skip_empty_area` to `False`.
         """
 
     def iter_rows(
@@ -89,7 +90,7 @@ class CalamineSheet:
         """Retunrning data from sheet as iterator of lists."""
 
 @typing.final
-class CalamineWorkbook:
+class CalamineWorkbook(contextlib.AbstractContextManager):
     path: str | None
     sheet_names: list[str]
     sheets_metadata: list[SheetMetadata]
@@ -99,49 +100,83 @@ class CalamineWorkbook:
     ) -> "CalamineWorkbook":
         """Determining type of pyobject and reading from it.
 
-        Parameters
-        ----------
-        path_or_filelike :
-            - path (string),
-            - pathlike (pathlib.Path),
-            - IO (must imlpement read/seek methods).
+        Args:
+            path_or_filelike (str | os.PathLike | ReadBuffer): path to file or IO (must imlpement read/seek methods).
         """
 
     @classmethod
     def from_path(cls, path: str | os.PathLike) -> "CalamineWorkbook":
         """Reading file from path.
 
-        Parameters
-        ----------
-        path : path (string)."""
+        Args:
+            path (str | os.PathLike): path to file.
+        """
 
     @classmethod
     def from_filelike(cls, filelike: ReadBuffer) -> "CalamineWorkbook":
         """Reading file from IO.
 
-        Parameters
-        ----------
-        filelike : IO (must imlpement read/seek methods).
+        Args:
+            filelike : IO (must imlpement read/seek methods).
+        """
+
+    def close(self) -> None:
+        """Close the workbook.
+
+        Drop internal rust structure from workbook (and close the file under the hood).
+        `get_sheet_by_name`/`get_sheet_by_index` will raise WorkbookClosed after calling that method.
+
+        Raises:
+            WorkbookClosed: If workbook already closed.
+        """
+
+    def __enter__(self) -> "CalamineWorkbook": ...
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: types.TracebackType | None,
+    ) -> None: ...
+    def get_sheet_by_name(self, name: str) -> CalamineSheet:
+        """Get worksheet by name.
+
+        Args:
+            name(str): name of worksheet
+
+        Returns:
+            CalamineSheet
+
+        Raises:
+            WorkbookClosed: If workbook already closed.
+            WorksheetNotFound: If worksheet not found in workbook.
         """
 
-    def get_sheet_by_name(self, name: str) -> CalamineSheet: ...
-    def get_sheet_by_index(self, index: int) -> CalamineSheet: ...
+    def get_sheet_by_index(self, index: int) -> CalamineSheet:
+        """Get worksheet by index.
+
+        Args:
+            index(int): index of worksheet
+
+        Returns:
+            CalamineSheet
+
+        Raises:
+            WorkbookClosed: If workbook already closed.
+            WorksheetNotFound: If worksheet not found in workbook.
+        """
 
 class CalamineError(Exception): ...
 class PasswordError(CalamineError): ...
 class WorksheetNotFound(CalamineError): ...
 class XmlError(CalamineError): ...
 class ZipError(CalamineError): ...
+class WorkbookClosed(CalamineError): ...
 
 def load_workbook(
     path_or_filelike: str | os.PathLike | ReadBuffer,
 ) -> CalamineWorkbook:
     """Determining type of pyobject and reading from it.
 
-    Parameters
-    ----------
-    path_or_filelike :
-        - path (string),
-        - pathlike (pathlib.Path),
-        - IO (must imlpement read/seek methods).
+    Args:
+        path_or_filelike (str | os.PathLike | ReadBuffer): path to file or IO (must imlpement read/seek methods).
     """
diff --git a/src/lib.rs b/src/lib.rs
index e8967b2..f01b4bc 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -3,8 +3,8 @@ use pyo3::prelude::*;
 mod types;
 mod utils;
 use crate::types::{
-    CalamineError, CalamineSheet, CalamineWorkbook, CellValue, PasswordError, SheetMetadata,
-    SheetTypeEnum, SheetVisibleEnum, WorksheetNotFound, XmlError, ZipError,
+    CalamineError, CalamineSheet, CalamineWorkbook, CellValue, Error, PasswordError, SheetMetadata,
+    SheetTypeEnum, SheetVisibleEnum, WorkbookClosed, WorksheetNotFound, XmlError, ZipError,
 };
 
 #[pyfunction]
@@ -28,5 +28,6 @@ fn _python_calamine(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     )?;
     m.add("XmlError", py.get_type_bound::<XmlError>())?;
     m.add("ZipError", py.get_type_bound::<ZipError>())?;
+    m.add("WorkbookClosed", py.get_type_bound::<WorkbookClosed>())?;
     Ok(())
 }
diff --git a/src/types/errors.rs b/src/types/errors.rs
new file mode 100644
index 0000000..a4220b9
--- /dev/null
+++ b/src/types/errors.rs
@@ -0,0 +1,16 @@
+use calamine::Error as CalamineCrateError;
+use pyo3::create_exception;
+use pyo3::exceptions::PyException;
+
+#[derive(Debug)]
+pub enum Error {
+    Calamine(CalamineCrateError),
+    WorkbookClosed,
+}
+
+create_exception!(python_calamine, CalamineError, PyException);
+create_exception!(python_calamine, PasswordError, CalamineError);
+create_exception!(python_calamine, WorksheetNotFound, CalamineError);
+create_exception!(python_calamine, XmlError, CalamineError);
+create_exception!(python_calamine, ZipError, CalamineError);
+create_exception!(python_calamine, WorkbookClosed, CalamineError);
diff --git a/src/types/mod.rs b/src/types/mod.rs
index dc66572..60fed03 100644
--- a/src/types/mod.rs
+++ b/src/types/mod.rs
@@ -1,15 +1,10 @@
-use pyo3::create_exception;
-use pyo3::exceptions::PyException;
-
 mod cell;
+mod errors;
 mod sheet;
 mod workbook;
 pub use cell::CellValue;
+pub use errors::{
+    CalamineError, Error, PasswordError, WorkbookClosed, WorksheetNotFound, XmlError, ZipError,
+};
 pub use sheet::{CalamineSheet, SheetMetadata, SheetTypeEnum, SheetVisibleEnum};
 pub use workbook::CalamineWorkbook;
-
-create_exception!(python_calamine, CalamineError, PyException);
-create_exception!(python_calamine, PasswordError, CalamineError);
-create_exception!(python_calamine, WorksheetNotFound, CalamineError);
-create_exception!(python_calamine, XmlError, CalamineError);
-create_exception!(python_calamine, ZipError, CalamineError);
diff --git a/src/types/workbook.rs b/src/types/workbook.rs
index 97a010b..6a0af10 100644
--- a/src/types/workbook.rs
+++ b/src/types/workbook.rs
@@ -2,18 +2,19 @@ use std::fs::File;
 use std::io::{BufReader, Cursor, Read};
 use std::path::PathBuf;
 
-use calamine::{open_workbook_auto, open_workbook_auto_from_rs, Error, Reader, Sheets};
+use calamine::{open_workbook_auto, open_workbook_auto_from_rs, Reader, Sheets};
 use pyo3::exceptions::PyTypeError;
 use pyo3::prelude::*;
 use pyo3::types::{PyString, PyType};
 use pyo3_file::PyFileLikeObject;
 
 use crate::utils::err_to_py;
-use crate::{CalamineSheet, SheetMetadata, WorksheetNotFound};
+use crate::{CalamineSheet, Error, SheetMetadata, WorksheetNotFound};
 
 enum SheetsEnum {
     File(Sheets<BufReader<File>>),
     FileLike(Sheets<Cursor<Vec<u8>>>),
+    None,
 }
 
 impl SheetsEnum {
@@ -21,6 +22,7 @@ impl SheetsEnum {
         match self {
             SheetsEnum::File(f) => f.sheets_metadata(),
             SheetsEnum::FileLike(f) => f.sheets_metadata(),
+            SheetsEnum::None => unreachable!(),
         }
         .iter()
         .map(|s| SheetMetadata::new(s.name.clone(), s.typ, s.visible))
@@ -31,13 +33,15 @@ impl SheetsEnum {
         match self {
             SheetsEnum::File(f) => f.sheet_names(),
             SheetsEnum::FileLike(f) => f.sheet_names(),
+            SheetsEnum::None => unreachable!(),
         }
     }
 
     fn worksheet_range(&mut self, name: &str) -> Result<calamine::Range<calamine::Data>, Error> {
         match self {
-            SheetsEnum::File(f) => f.worksheet_range(name),
-            SheetsEnum::FileLike(f) => f.worksheet_range(name),
+            SheetsEnum::File(f) => f.worksheet_range(name).map_err(Error::Calamine),
+            SheetsEnum::FileLike(f) => f.worksheet_range(name).map_err(Error::Calamine),
+            SheetsEnum::None => Err(Error::WorkbookClosed),
         }
     }
 }
@@ -107,6 +111,30 @@ impl CalamineWorkbook {
     fn py_get_sheet_by_index(&mut self, py: Python<'_>, index: usize) -> PyResult<CalamineSheet> {
         py.allow_threads(|| self.get_sheet_by_index(index))
     }
+
+    fn close(&mut self) -> PyResult<()> {
+        match self.sheets {
+            SheetsEnum::None => Err(Error::WorkbookClosed),
+            _ => {
+                self.sheets = SheetsEnum::None;
+                Ok(())
+            }
+        }
+        .map_err(err_to_py)
+    }
+
+    fn __enter__(slf: Py<Self>) -> Py<Self> {
+        slf
+    }
+
+    fn __exit__(
+        &mut self,
+        _exc_type: PyObject,
+        _exc_value: PyObject,
+        _traceback: PyObject,
+    ) -> PyResult<()> {
+        self.close()
+    }
 }
 
 impl CalamineWorkbook {
@@ -129,7 +157,11 @@ impl CalamineWorkbook {
         PyFileLikeObject::with_requirements(filelike, true, false, true, false)?
             .read_to_end(&mut buf)?;
         let reader = Cursor::new(buf);
-        let sheets = SheetsEnum::FileLike(open_workbook_auto_from_rs(reader).map_err(err_to_py)?);
+        let sheets = SheetsEnum::FileLike(
+            open_workbook_auto_from_rs(reader)
+                .map_err(Error::Calamine)
+                .map_err(err_to_py)?,
+        );
         let sheet_names = sheets.sheet_names().to_owned();
         let sheets_metadata = sheets.sheets_metadata().to_owned();
 
@@ -142,7 +174,11 @@ impl CalamineWorkbook {
     }
 
     pub fn from_path(path: &str) -> PyResult<Self> {
-        let sheets = SheetsEnum::File(open_workbook_auto(path).map_err(err_to_py)?);
+        let sheets = SheetsEnum::File(
+            open_workbook_auto(path)
+                .map_err(Error::Calamine)
+                .map_err(err_to_py)?,
+        );
         let sheet_names = sheets.sheet_names().to_owned();
         let sheets_metadata = sheets.sheets_metadata().to_owned();
 
diff --git a/src/utils.rs b/src/utils.rs
index f79bd0a..185c33d 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -1,46 +1,55 @@
-use calamine::{Error, OdsError, XlsError, XlsbError, XlsxError};
+use calamine::{Error as CalamineCrateError, OdsError, XlsError, XlsbError, XlsxError};
 use pyo3::exceptions::PyIOError;
 use pyo3::PyErr;
 
-use crate::{CalamineError, PasswordError, WorksheetNotFound, XmlError, ZipError};
+use crate::{
+    CalamineError, Error, PasswordError, WorkbookClosed, WorksheetNotFound, XmlError, ZipError,
+};
 
 pub fn err_to_py(e: Error) -> PyErr {
     match e {
-        Error::Io(err) => PyIOError::new_err(err.to_string()),
-        Error::Ods(ref err) => match err {
-            OdsError::Io(error) => PyIOError::new_err(error.to_string()),
-            OdsError::Zip(error) => ZipError::new_err(error.to_string()),
-            OdsError::Xml(error) => XmlError::new_err(error.to_string()),
-            OdsError::XmlAttr(error) => XmlError::new_err(error.to_string()),
-            OdsError::Password => PasswordError::new_err(err.to_string()),
-            OdsError::WorksheetNotFound(error) => WorksheetNotFound::new_err(error.to_string()),
-            _ => CalamineError::new_err(err.to_string()),
+        Error::Calamine(calamine_error) => match calamine_error {
+            CalamineCrateError::Io(err) => PyIOError::new_err(err.to_string()),
+            CalamineCrateError::Ods(ref err) => match err {
+                OdsError::Io(error) => PyIOError::new_err(error.to_string()),
+                OdsError::Zip(error) => ZipError::new_err(error.to_string()),
+                OdsError::Xml(error) => XmlError::new_err(error.to_string()),
+                OdsError::XmlAttr(error) => XmlError::new_err(error.to_string()),
+                OdsError::Password => PasswordError::new_err(err.to_string()),
+                OdsError::WorksheetNotFound(error) => WorksheetNotFound::new_err(error.to_string()),
+                _ => CalamineError::new_err(err.to_string()),
+            },
+            CalamineCrateError::Xls(ref err) => match err {
+                XlsError::Io(error) => PyIOError::new_err(error.to_string()),
+                XlsError::Password => PasswordError::new_err(err.to_string()),
+                XlsError::WorksheetNotFound(error) => WorksheetNotFound::new_err(error.to_string()),
+                _ => CalamineError::new_err(err.to_string()),
+            },
+            CalamineCrateError::Xlsx(ref err) => match err {
+                XlsxError::Io(error) => PyIOError::new_err(error.to_string()),
+                XlsxError::Zip(error) => ZipError::new_err(error.to_string()),
+                XlsxError::Xml(error) => XmlError::new_err(error.to_string()),
+                XlsxError::XmlAttr(error) => XmlError::new_err(error.to_string()),
+                XlsxError::XmlEof(error) => XmlError::new_err(error.to_string()),
+                XlsxError::Password => PasswordError::new_err(err.to_string()),
+                XlsxError::WorksheetNotFound(error) => {
+                    WorksheetNotFound::new_err(error.to_string())
+                }
+                _ => CalamineError::new_err(err.to_string()),
+            },
+            CalamineCrateError::Xlsb(ref err) => match err {
+                XlsbError::Io(error) => PyIOError::new_err(error.to_string()),
+                XlsbError::Zip(error) => ZipError::new_err(error.to_string()),
+                XlsbError::Xml(error) => XmlError::new_err(error.to_string()),
+                XlsbError::XmlAttr(error) => XmlError::new_err(error.to_string()),
+                XlsbError::Password => PasswordError::new_err(err.to_string()),
+                XlsbError::WorksheetNotFound(error) => {
+                    WorksheetNotFound::new_err(error.to_string())
+                }
+                _ => CalamineError::new_err(err.to_string()),
+            },
+            _ => CalamineError::new_err(calamine_error.to_string()),
         },
-        Error::Xls(ref err) => match err {
-            XlsError::Io(error) => PyIOError::new_err(error.to_string()),
-            XlsError::Password => PasswordError::new_err(err.to_string()),
-            XlsError::WorksheetNotFound(error) => WorksheetNotFound::new_err(error.to_string()),
-            _ => CalamineError::new_err(err.to_string()),
-        },
-        Error::Xlsx(ref err) => match err {
-            XlsxError::Io(error) => PyIOError::new_err(error.to_string()),
-            XlsxError::Zip(error) => ZipError::new_err(error.to_string()),
-            XlsxError::Xml(error) => XmlError::new_err(error.to_string()),
-            XlsxError::XmlAttr(error) => XmlError::new_err(error.to_string()),
-            XlsxError::XmlEof(error) => XmlError::new_err(error.to_string()),
-            XlsxError::Password => PasswordError::new_err(err.to_string()),
-            XlsxError::WorksheetNotFound(error) => WorksheetNotFound::new_err(error.to_string()),
-            _ => CalamineError::new_err(err.to_string()),
-        },
-        Error::Xlsb(ref err) => match err {
-            XlsbError::Io(error) => PyIOError::new_err(error.to_string()),
-            XlsbError::Zip(error) => ZipError::new_err(error.to_string()),
-            XlsbError::Xml(error) => XmlError::new_err(error.to_string()),
-            XlsbError::XmlAttr(error) => XmlError::new_err(error.to_string()),
-            XlsbError::Password => PasswordError::new_err(err.to_string()),
-            XlsbError::WorksheetNotFound(error) => WorksheetNotFound::new_err(error.to_string()),
-            _ => CalamineError::new_err(err.to_string()),
-        },
-        _ => CalamineError::new_err(e.to_string()),
+        Error::WorkbookClosed => WorkbookClosed::new_err("".to_string()),
     }
 }


Test patch:
diff --git a/tests/test_base.py b/tests/test_base.py
index 17e9cc4..2cbd49f 100644
--- a/tests/test_base.py
+++ b/tests/test_base.py
@@ -3,7 +3,13 @@
 from pathlib import Path
 
 import pytest
-from python_calamine import CalamineWorkbook, PasswordError, WorksheetNotFound, ZipError
+from python_calamine import (
+    CalamineWorkbook,
+    PasswordError,
+    WorkbookClosed,
+    WorksheetNotFound,
+    ZipError,
+)
 
 PATH = Path(__file__).parent / "data"
 
@@ -341,3 +347,19 @@ def test_path_or_filelike(obj):
 def test_path_or_filelike_error():
     with pytest.raises(TypeError):
         CalamineWorkbook.from_object(object())
+
+
+def test_close_workbook():
+    reader = CalamineWorkbook.from_path(PATH / "base.xlsx")
+    reader.close()
+
+    with pytest.raises(WorkbookClosed):
+        reader.get_sheet_by_index(1)
+
+
+def test_close_workbook_double():
+    reader = CalamineWorkbook.from_path(PATH / "base.xlsx")
+    reader.close()
+
+    with pytest.raises(WorkbookClosed):
+        reader.close()