Repository: pola-rs/polars
File: py-polars/tests/unit/io/test_spreadsheet.py (Python)

Code before:
    df = pl.DataFrame({"colx": ["aaa", "bbb", "ccc"], "coly": [-1234, 0, 5678]})

    with Workbook(xls := BytesIO()) as wb:
        ws = wb.add_worksheet("frame_data")
        df.write_excel(wb, worksheet=ws)
        ws.hide_zero()

    assert_frame_equal(df, pl.read_excel(xls, sheet_name="frame_data"))

    with pytest.raises(  # noqa: SIM117
        ValueError,
        match="the given workbook object .* is not the parent of worksheet 'frame_data'",
    ):
        with Workbook(BytesIO()) as wb:
            df.write_excel(wb, worksheet=ws)

    with pytest.raises(  # noqa: SIM117
        TypeError,
        match="worksheet object requires the parent workbook object; found workbook=None",
    ):
        with Workbook(BytesIO()) as wb:
            df.write_excel(None, worksheet=ws)


def test_excel_write_beyond_max_rows_cols() -> None:
    path = "/tmp/test.xlsx"
    sheet = "mysheet"

    df = pl.DataFrame({"col1": range(10), "col2": range(10, 20)})

    with pytest.raises(pl.exceptions.InvalidOperationError):
        df.write_excel(workbook=path, worksheet=sheet, position="A1048570")


def test_excel_freeze_panes() -> None:
    from xlsxwriter import Workbook

    # note: checks that empty tables don't error on write
    df1 = pl.DataFrame(schema={"colx": pl.Date, "coly": pl.String, "colz": pl.Float64})
    df2 = pl.DataFrame(schema={"colx": pl.Date, "coly": pl.String, "colz": pl.Float64})
    df3 = pl.DataFrame(schema={"colx": pl.Date, "coly": pl.String, "colz": pl.Float64})

    xls = BytesIO()

    # use all three freeze_pane notations
    with Workbook(xls) as wb:
        df1.write_excel(workbook=wb, worksheet="sheet1", freeze_panes=(1, 0))
        df2.write_excel(workbook=wb, worksheet="sheet2", freeze_panes=(1, 0, 3, 4))
        df3.write_excel(workbook=wb, worksheet="sheet3", freeze_panes=("B2"))

    table_names: set[str] = set()

Review comment (suggestion):
```suggestion
def test_excel_write_beyond_max_rows_cols(tmp_path: Path) -> None:
    tmp_path.mkdir(exist_ok=True)
    path = tmp_path / "test_max_dimensions.xlsx"
```

Code after:
    df = pl.DataFrame({"colx": ["aaa", "bbb", "ccc"], "coly": [-1234, 0, 5678]})

    with Workbook(xls := BytesIO()) as wb:
        ws = wb.add_worksheet("frame_data")
        df.write_excel(wb, worksheet=ws)
        ws.hide_zero()

    assert_frame_equal(df, pl.read_excel(xls, sheet_name="frame_data"))

    with pytest.raises(  # noqa: SIM117
        ValueError,
        match="the given workbook object .* is not the parent of worksheet 'frame_data'",
    ):
        with Workbook(BytesIO()) as wb:
            df.write_excel(wb, worksheet=ws)

    with pytest.raises(  # noqa: SIM117
        TypeError,
        match="worksheet object requires the parent workbook object; found workbook=None",
    ):
        with Workbook(BytesIO()) as wb:
            df.write_excel(None, worksheet=ws)


def test_excel_write_beyond_max_rows_cols(tmp_path: Path) -> None:
    tmp_path.mkdir(exist_ok=True)
    path = tmp_path / "test_max_dimensions.xlsx"
    sheet = "mysheet"

    df = pl.DataFrame({"col1": range(10), "col2": range(10, 20)})

    with pytest.raises(pl.exceptions.InvalidOperationError):
        df.write_excel(workbook=path, worksheet=sheet, position="A1048570")


def test_excel_freeze_panes() -> None:
    from xlsxwriter import Workbook

    # note: checks that empty tables don't error on write
    df1 = pl.DataFrame(schema={"colx": pl.Date, "coly": pl.String, "colz": pl.Float64})
    df2 = pl.DataFrame(schema={"colx": pl.Date, "coly": pl.String, "colz": pl.Float64})
    df3 = pl.DataFrame(schema={"colx": pl.Date, "coly": pl.String, "colz": pl.Float64})

    xls = BytesIO()

    # use all three freeze_pane notations
    with Workbook(xls) as wb:
        df1.write_excel(workbook=wb, worksheet="sheet1", freeze_panes=(1, 0))
        df2.write_excel(workbook=wb, worksheet="sheet2", freeze_panes=(1, 0, 3, 4))
        df3.write_excel(workbook=wb, worksheet="sheet3", freeze_panes=("B2"))