Skip to content

Commit f0c6b59

Browse files
authored
ENH: Allow safe access to .book in ExcelWriter (#45687)
1 parent 6781480 commit f0c6b59

10 files changed

+82
-10
lines changed

pandas/io/excel/_base.py

+6-1
Original file line numberDiff line numberDiff line change
@@ -1048,6 +1048,12 @@ def engine(self) -> str:
10481048
"""Name of engine."""
10491049
pass
10501050

1051+
@property
1052+
@abc.abstractmethod
1053+
def sheets(self) -> dict[str, Any]:
1054+
"""Mapping of sheet names to sheet objects."""
1055+
pass
1056+
10511057
@abc.abstractmethod
10521058
def write_cells(
10531059
self,
@@ -1112,7 +1118,6 @@ def __init__(
11121118
self.handles = get_handle(
11131119
path, mode, storage_options=storage_options, is_text=False
11141120
)
1115-
self.sheets: dict[str, Any] = {}
11161121
self.cur_sheet = None
11171122

11181123
if date_format is None:

pandas/io/excel/_odswriter.py

+12-1
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,17 @@ def __init__(
5858
self.book = OpenDocumentSpreadsheet(**engine_kwargs)
5959
self._style_dict: dict[str, str] = {}
6060

61+
@property
62+
def sheets(self) -> dict[str, Any]:
63+
"""Mapping of sheet names to sheet objects."""
64+
from odf.table import Table
65+
66+
result = {
67+
sheet.getAttribute("name"): sheet
68+
for sheet in self.book.getElementsByType(Table)
69+
}
70+
return result
71+
6172
def save(self) -> None:
6273
"""
6374
Save workbook to disk.
@@ -91,7 +102,7 @@ def write_cells(
91102
wks = self.sheets[sheet_name]
92103
else:
93104
wks = Table(name=sheet_name)
94-
self.sheets[sheet_name] = wks
105+
self.book.spreadsheet.addElement(wks)
95106

96107
if validate_freeze_panes(freeze_panes):
97108
freeze_panes = cast(Tuple[int, int], freeze_panes)

pandas/io/excel/_openpyxl.py

+6-4
Original file line numberDiff line numberDiff line change
@@ -68,15 +68,19 @@ def __init__(
6868

6969
self.book = load_workbook(self.handles.handle, **engine_kwargs)
7070
self.handles.handle.seek(0)
71-
self.sheets = {name: self.book[name] for name in self.book.sheetnames}
72-
7371
else:
7472
# Create workbook object with default optimized_write=True.
7573
self.book = Workbook(**engine_kwargs)
7674

7775
if self.book.worksheets:
7876
self.book.remove(self.book.worksheets[0])
7977

78+
@property
79+
def sheets(self) -> dict[str, Any]:
80+
"""Mapping of sheet names to sheet objects."""
81+
result = {name: self.book[name] for name in self.book.sheetnames}
82+
return result
83+
8084
def save(self) -> None:
8185
"""
8286
Save workbook to disk.
@@ -440,7 +444,6 @@ def write_cells(
440444
target_index = self.book.index(old_wks)
441445
del self.book[sheet_name]
442446
wks = self.book.create_sheet(sheet_name, target_index)
443-
self.sheets[sheet_name] = wks
444447
elif self.if_sheet_exists == "error":
445448
raise ValueError(
446449
f"Sheet '{sheet_name}' already exists and "
@@ -458,7 +461,6 @@ def write_cells(
458461
else:
459462
wks = self.book.create_sheet()
460463
wks.title = sheet_name
461-
self.sheets[sheet_name] = wks
462464

463465
if validate_freeze_panes(freeze_panes):
464466
freeze_panes = cast(Tuple[int, int], freeze_panes)

pandas/io/excel/_xlsxwriter.py

+7-4
Original file line numberDiff line numberDiff line change
@@ -205,6 +205,11 @@ def __init__(
205205

206206
self.book = Workbook(self.handles.handle, **engine_kwargs)
207207

208+
@property
209+
def sheets(self) -> dict[str, Any]:
210+
result = self.book.sheetnames
211+
return result
212+
208213
def save(self) -> None:
209214
"""
210215
Save workbook to disk.
@@ -222,11 +227,9 @@ def write_cells(
222227
# Write the frame cells using xlsxwriter.
223228
sheet_name = self._get_sheet_name(sheet_name)
224229

225-
if sheet_name in self.sheets:
226-
wks = self.sheets[sheet_name]
227-
else:
230+
wks = self.book.get_worksheet_by_name(sheet_name)
231+
if wks is None:
228232
wks = self.book.add_worksheet(sheet_name)
229-
self.sheets[sheet_name] = wks
230233

231234
style_dict = {"null": None}
232235

pandas/io/excel/_xlwt.py

+6
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,12 @@ def __init__(
6363
self.fm_datetime = xlwt.easyxf(num_format_str=self.datetime_format)
6464
self.fm_date = xlwt.easyxf(num_format_str=self.date_format)
6565

66+
@property
67+
def sheets(self) -> dict[str, Any]:
68+
"""Mapping of sheet names to sheet objects."""
69+
result = {sheet.name: sheet for sheet in self.book._Workbook__worksheets}
70+
return result
71+
6672
def save(self) -> None:
6773
"""
6874
Save workbook to disk.

pandas/tests/io/excel/test_odswriter.py

+10
Original file line numberDiff line numberDiff line change
@@ -56,3 +56,13 @@ def test_engine_kwargs(ext, engine_kwargs):
5656
else:
5757
with ExcelWriter(f, engine="odf", engine_kwargs=engine_kwargs) as _:
5858
pass
59+
60+
61+
def test_book_and_sheets_consistent(ext):
62+
# GH#45687 - Ensure sheets is updated if user modifies book
63+
with tm.ensure_clean(ext) as f:
64+
with ExcelWriter(f) as writer:
65+
assert writer.sheets == {}
66+
table = odf.table.Table(name="test_name")
67+
writer.book.spreadsheet.addElement(table)
68+
assert writer.sheets == {"test_name": table}

pandas/tests/io/excel/test_openpyxl.py

+9
Original file line numberDiff line numberDiff line change
@@ -379,3 +379,12 @@ def test_read_empty_with_blank_row(datapath, ext, read_only):
379379
result = pd.read_excel(wb, engine="openpyxl")
380380
expected = DataFrame()
381381
tm.assert_frame_equal(result, expected)
382+
383+
384+
def test_book_and_sheets_consistent(ext):
385+
# GH#45687 - Ensure sheets is updated if user modifies book
386+
with tm.ensure_clean(ext) as f:
387+
with ExcelWriter(f, engine="openpyxl") as writer:
388+
assert writer.sheets == {}
389+
sheet = writer.book.create_sheet("test_name", 0)
390+
assert writer.sheets == {"test_name": sheet}

pandas/tests/io/excel/test_writers.py

+8
Original file line numberDiff line numberDiff line change
@@ -1271,10 +1271,12 @@ def test_register_writer(self):
12711271
# some awkward mocking to test out dispatch and such actually works
12721272
called_save = []
12731273
called_write_cells = []
1274+
called_sheets = []
12741275

12751276
class DummyClass(ExcelWriter):
12761277
called_save = False
12771278
called_write_cells = False
1279+
called_sheets = False
12781280
supported_extensions = ["xlsx", "xls"]
12791281
engine = "dummy"
12801282

@@ -1284,12 +1286,18 @@ def save(self):
12841286
def write_cells(self, *args, **kwargs):
12851287
called_write_cells.append(True)
12861288

1289+
@property
1290+
def sheets(self):
1291+
called_sheets.append(True)
1292+
12871293
def check_called(func):
12881294
func()
12891295
assert len(called_save) >= 1
12901296
assert len(called_write_cells) >= 1
1297+
assert len(called_sheets) == 0
12911298
del called_save[:]
12921299
del called_write_cells[:]
1300+
del called_sheets[:]
12931301

12941302
with option_context("io.excel.xlsx.writer", "dummy"):
12951303
path = "something.xlsx"

pandas/tests/io/excel/test_xlsxwriter.py

+9
Original file line numberDiff line numberDiff line change
@@ -83,3 +83,12 @@ def test_engine_kwargs(ext, nan_inf_to_errors):
8383
with tm.ensure_clean(ext) as f:
8484
with ExcelWriter(f, engine="xlsxwriter", engine_kwargs=engine_kwargs) as writer:
8585
assert writer.book.nan_inf_to_errors == nan_inf_to_errors
86+
87+
88+
def test_book_and_sheets_consistent(ext):
89+
# GH#45687 - Ensure sheets is updated if user modifies book
90+
with tm.ensure_clean(ext) as f:
91+
with ExcelWriter(f, engine="xlsxwriter") as writer:
92+
assert writer.sheets == {}
93+
sheet = writer.book.add_worksheet("test_name")
94+
assert writer.sheets == {"test_name": sheet}

pandas/tests/io/excel/test_xlwt.py

+9
Original file line numberDiff line numberDiff line change
@@ -125,3 +125,12 @@ def test_engine_kwargs(ext, style_compression):
125125
assert writer.book._Workbook__styles.style_compression == style_compression
126126
# xlwt won't allow us to close without writing something
127127
DataFrame().to_excel(writer)
128+
129+
130+
def test_book_and_sheets_consistent(ext):
131+
# GH#45687 - Ensure sheets is updated if user modifies book
132+
with tm.ensure_clean(ext) as f:
133+
with ExcelWriter(f) as writer:
134+
assert writer.sheets == {}
135+
sheet = writer.book.add_sheet("test_name")
136+
assert writer.sheets == {"test_name": sheet}

0 commit comments

Comments
 (0)