ENH: Allow safe access to .book in ExcelWriter (pandas-dev#45687)

Author: rhshadrach

pandas/io/excel/_base.py

@@ -1048,6 +1048,12 @@ def engine(self) -> str:
         """Name of engine."""
         pass
 
+    @property
+    @abc.abstractmethod
+    def sheets(self) -> dict[str, Any]:
+        """Mapping of sheet names to sheet objects."""
+        pass
+
     @abc.abstractmethod
     def write_cells(
         self,
@@ -1112,7 +1118,6 @@ def __init__(
             self.handles = get_handle(
                 path, mode, storage_options=storage_options, is_text=False
             )
-        self.sheets: dict[str, Any] = {}
         self.cur_sheet = None
 
         if date_format is None:

pandas/io/excel/_odswriter.py

@@ -58,6 +58,17 @@ def __init__(
         self.book = OpenDocumentSpreadsheet(**engine_kwargs)
         self._style_dict: dict[str, str] = {}
 
+    @property
+    def sheets(self) -> dict[str, Any]:
+        """Mapping of sheet names to sheet objects."""
+        from odf.table import Table
+
+        result = {
+            sheet.getAttribute("name"): sheet
+            for sheet in self.book.getElementsByType(Table)
+        }
+        return result
+
     def save(self) -> None:
         """
         Save workbook to disk.
@@ -91,7 +102,7 @@ def write_cells(
             wks = self.sheets[sheet_name]
         else:
             wks = Table(name=sheet_name)
-            self.sheets[sheet_name] = wks
+            self.book.spreadsheet.addElement(wks)
 
         if validate_freeze_panes(freeze_panes):
             freeze_panes = cast(Tuple[int, int], freeze_panes)

pandas/io/excel/_openpyxl.py

@@ -68,15 +68,19 @@ def __init__(
 
             self.book = load_workbook(self.handles.handle, **engine_kwargs)
             self.handles.handle.seek(0)
-            self.sheets = {name: self.book[name] for name in self.book.sheetnames}
-
         else:
             # Create workbook object with default optimized_write=True.
             self.book = Workbook(**engine_kwargs)
 
             if self.book.worksheets:
                 self.book.remove(self.book.worksheets[0])
 
+    @property
+    def sheets(self) -> dict[str, Any]:
+        """Mapping of sheet names to sheet objects."""
+        result = {name: self.book[name] for name in self.book.sheetnames}
+        return result
+
     def save(self) -> None:
         """
         Save workbook to disk.
@@ -440,7 +444,6 @@ def write_cells(
                     target_index = self.book.index(old_wks)
                     del self.book[sheet_name]
                     wks = self.book.create_sheet(sheet_name, target_index)
-                    self.sheets[sheet_name] = wks
                 elif self.if_sheet_exists == "error":
                     raise ValueError(
                         f"Sheet '{sheet_name}' already exists and "
@@ -458,7 +461,6 @@ def write_cells(
         else:
             wks = self.book.create_sheet()
             wks.title = sheet_name
-            self.sheets[sheet_name] = wks
 
         if validate_freeze_panes(freeze_panes):
             freeze_panes = cast(Tuple[int, int], freeze_panes)

pandas/io/excel/_xlsxwriter.py

@@ -205,6 +205,11 @@ def __init__(
 
         self.book = Workbook(self.handles.handle, **engine_kwargs)
 
+    @property
+    def sheets(self) -> dict[str, Any]:
+        result = self.book.sheetnames
+        return result
+
     def save(self) -> None:
         """
         Save workbook to disk.
@@ -222,11 +227,9 @@ def write_cells(
         # Write the frame cells using xlsxwriter.
         sheet_name = self._get_sheet_name(sheet_name)
 
-        if sheet_name in self.sheets:
-            wks = self.sheets[sheet_name]
-        else:
+        wks = self.book.get_worksheet_by_name(sheet_name)
+        if wks is None:
             wks = self.book.add_worksheet(sheet_name)
-            self.sheets[sheet_name] = wks
 
         style_dict = {"null": None}
 

pandas/io/excel/_xlwt.py

@@ -63,6 +63,12 @@ def __init__(
         self.fm_datetime = xlwt.easyxf(num_format_str=self.datetime_format)
         self.fm_date = xlwt.easyxf(num_format_str=self.date_format)
 
+    @property
+    def sheets(self) -> dict[str, Any]:
+        """Mapping of sheet names to sheet objects."""
+        result = {sheet.name: sheet for sheet in self.book._Workbook__worksheets}
+        return result
+
     def save(self) -> None:
         """
         Save workbook to disk.

pandas/tests/io/excel/test_odswriter.py

@@ -56,3 +56,13 @@ def test_engine_kwargs(ext, engine_kwargs):
         else:
             with ExcelWriter(f, engine="odf", engine_kwargs=engine_kwargs) as _:
                 pass
+
+
+def test_book_and_sheets_consistent(ext):
+    # GH#45687 - Ensure sheets is updated if user modifies book
+    with tm.ensure_clean(ext) as f:
+        with ExcelWriter(f) as writer:
+            assert writer.sheets == {}
+            table = odf.table.Table(name="test_name")
+            writer.book.spreadsheet.addElement(table)
+            assert writer.sheets == {"test_name": table}

pandas/tests/io/excel/test_openpyxl.py

@@ -379,3 +379,12 @@ def test_read_empty_with_blank_row(datapath, ext, read_only):
             result = pd.read_excel(wb, engine="openpyxl")
     expected = DataFrame()
     tm.assert_frame_equal(result, expected)
+
+
+def test_book_and_sheets_consistent(ext):
+    # GH#45687 - Ensure sheets is updated if user modifies book
+    with tm.ensure_clean(ext) as f:
+        with ExcelWriter(f, engine="openpyxl") as writer:
+            assert writer.sheets == {}
+            sheet = writer.book.create_sheet("test_name", 0)
+            assert writer.sheets == {"test_name": sheet}

pandas/tests/io/excel/test_writers.py

@@ -1271,10 +1271,12 @@ def test_register_writer(self):
         # some awkward mocking to test out dispatch and such actually works
         called_save = []
         called_write_cells = []
+        called_sheets = []
 
         class DummyClass(ExcelWriter):
             called_save = False
             called_write_cells = False
+            called_sheets = False
             supported_extensions = ["xlsx", "xls"]
             engine = "dummy"
 
@@ -1284,12 +1286,18 @@ def save(self):
             def write_cells(self, *args, **kwargs):
                 called_write_cells.append(True)
 
+            @property
+            def sheets(self):
+                called_sheets.append(True)
+
         def check_called(func):
             func()
             assert len(called_save) >= 1
             assert len(called_write_cells) >= 1
+            assert len(called_sheets) == 0
             del called_save[:]
             del called_write_cells[:]
+            del called_sheets[:]
 
         with option_context("io.excel.xlsx.writer", "dummy"):
             path = "something.xlsx"

pandas/tests/io/excel/test_xlsxwriter.py

@@ -83,3 +83,12 @@ def test_engine_kwargs(ext, nan_inf_to_errors):
     with tm.ensure_clean(ext) as f:
         with ExcelWriter(f, engine="xlsxwriter", engine_kwargs=engine_kwargs) as writer:
             assert writer.book.nan_inf_to_errors == nan_inf_to_errors
+
+
+def test_book_and_sheets_consistent(ext):
+    # GH#45687 - Ensure sheets is updated if user modifies book
+    with tm.ensure_clean(ext) as f:
+        with ExcelWriter(f, engine="xlsxwriter") as writer:
+            assert writer.sheets == {}
+            sheet = writer.book.add_worksheet("test_name")
+            assert writer.sheets == {"test_name": sheet}

pandas/tests/io/excel/test_xlwt.py

@@ -125,3 +125,12 @@ def test_engine_kwargs(ext, style_compression):
             assert writer.book._Workbook__styles.style_compression == style_compression
             # xlwt won't allow us to close without writing something
             DataFrame().to_excel(writer)
+
+
+def test_book_and_sheets_consistent(ext):
+    # GH#45687 - Ensure sheets is updated if user modifies book
+    with tm.ensure_clean(ext) as f:
+        with ExcelWriter(f) as writer:
+            assert writer.sheets == {}
+            sheet = writer.book.add_sheet("test_name")
+            assert writer.sheets == {"test_name": sheet}