Skip to content

Commit 10d19d7

Browse files
committed
refactor: use type casting for Polars and Pandas DataFrames in dataframe2list function
1 parent f8607f9 commit 10d19d7

File tree

1 file changed

+4
-3
lines changed

1 file changed

+4
-3
lines changed

src/tppt/_pptx/shape/table.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
from pptx.enum.text import MSO_VERTICAL_ANCHOR, PP_ALIGN
66
from pptx.shapes.graphfrm import GraphicFrame
7+
78
from tppt._features import (
89
USE_PANDAS,
910
USE_POLARS,
@@ -152,19 +153,19 @@ def dataframe2list(data: DataFrame) -> list[list[str]]:
152153
if USE_POLARS:
153154
if isinstance(data, PolarsLazyFrame):
154155
# For LazyFrame, collect it first
155-
polars_df = data.collect()
156+
polars_df = cast(PolarsLazyFrame, data).collect()
156157
columns = list(polars_df.columns)
157158
rows = polars_df.to_numpy().tolist()
158159
return [columns] + rows
159160
elif isinstance(data, PolarsDataFrame):
160-
polars_df = data
161+
polars_df = cast(PolarsDataFrame, data)
161162
columns = list(polars_df.columns)
162163
rows = polars_df.to_numpy().tolist()
163164
return [columns] + rows
164165

165166
if USE_PANDAS and isinstance(data, PandasDataFrame): # type: ignore
166167
# Convert pandas DataFrame to list of lists
167-
pandas_df = data
168+
pandas_df = cast(PandasDataFrame, data)
168169
columns = pandas_df.columns.tolist()
169170
rows = pandas_df.values.tolist()
170171
return [columns] + rows

0 commit comments

Comments
 (0)