feat: enhance table data handling to support PolarsLazyFrame and add dataframe2list utility function

ya7010 · ya7010 · commit 05ddfc93979a · 2025-04-26T02:36:07.000+09:00
diff --git a/src/pptxr/_pptx/shape/table.py b/src/pptxr/_pptx/shape/table.py
@@ -18,7 +18,9 @@
 from . import Shape
 
 # Define DataFrame type alias
-DataFrame: TypeAlias = list[list[str]] | PandasDataFrame | PolarsDataFrame
+DataFrame: TypeAlias = (
+    list[list[str]] | PandasDataFrame | PolarsDataFrame | PolarsLazyFrame
+)
 
 
 class TableCellStyle(TypedDict):
@@ -48,7 +50,7 @@ class TableData(TableProps):
 
     type: Literal["table"]
 
-    data: DataFrame
+    data: list[list[str]]
 
 
 class Table(Shape[GraphicFrame]):
@@ -68,63 +70,46 @@ def __init__(
         table = pptx_obj.table
 
         # Apply first row as header if specified
-        if props.get("first_row_header"):
-            table.first_row = True
+        if (first_row := props.get("first_row_header")) is not None:
+            table.first_row = first_row
 
         # Apply table data if provided
-        if "data" in props:
-            data = props["data"]
-
-            # Convert different DataFrame types to list of lists
-            if USE_POLARS and issubclass(
-                data.__class__, (PolarsDataFrame, PolarsLazyFrame)
-            ):
-                # Convert pandas DataFrame to list of lists
-                data = cast(PandasDataFrame, data)
-                table_data = [data.columns.tolist()] + data.values.tolist()
-            elif USE_PANDAS and issubclass(data.__class__, PandasDataFrame):  # type: ignore
-                # Convert polars DataFrame to list of lists
-                data = cast(PandasDataFrame, data)
-                table_data = [data.columns] + data.to_numpy().tolist()
-            else:
-                table_data = data
-
+        if data := props.get("data"):
             # Now apply the data to the table
-            for i, row in enumerate(table_data):
+            for i, row in enumerate(data):
                 if i < len(table.rows):
                     for j, cell_text in enumerate(row):
                         if j < len(table.columns):
                             table.cell(i, j).text = str(cell_text)
 
         # Apply cell styles if provided
-        if "cell_styles" in props:
-            cell_styles = props["cell_styles"]
+        if (cell_styles := props.get("cell_styles")) is not None:
             for i, row_styles in enumerate(cell_styles):
                 if i < len(table.rows):
                     for j, cell_style in enumerate(row_styles):
                         if j < len(table.columns):
                             cell = table.cell(i, j)
 
-                            if "text_align" in cell_style:
+                            if (text_align := cell_style.get("text_align")) is not None:
                                 align_map = {
                                     "left": PP_ALIGN.LEFT,
                                     "center": PP_ALIGN.CENTER,
                                     "right": PP_ALIGN.RIGHT,
                                     "justify": PP_ALIGN.JUSTIFY,
                                 }
                                 paragraph = cell.text_frame.paragraphs[0]
-                                paragraph.alignment = align_map[
-                                    cell_style["text_align"]
-                                ]
+                                paragraph.alignment = align_map[text_align]
 
-                            if "vertical_align" in cell_style:
+                            if (
+                                vertical_align := cell_style.get("vertical_align")
+                            ) is not None:
                                 valign_map = {
                                     "top": MSO_VERTICAL_ANCHOR.TOP,
                                     "middle": MSO_VERTICAL_ANCHOR.MIDDLE,
                                     "bottom": MSO_VERTICAL_ANCHOR.BOTTOM,
                                 }
                                 cell.text_frame.vertical_anchor = valign_map[
-                                    cell_style["vertical_align"]
+                                    vertical_align
                                 ]
 
                             # Apply text formatting
@@ -161,3 +146,29 @@ def to_pptx(self) -> GraphicFrame:
     def from_pptx(cls, pptx_obj: GraphicFrame) -> Self:
         """Create from pptx table frame."""
         return cls(pptx_obj)
+
+
+def dataframe2list(data: DataFrame) -> list[list[str]]:
+    """Convert different DataFrame types to list of lists."""
+    if USE_POLARS:
+        if isinstance(data, PolarsLazyFrame):
+            # For LazyFrame, collect it first
+            polars_df = data.collect()
+            columns = list(polars_df.columns)
+            rows = polars_df.to_numpy().tolist()
+            return [columns] + rows
+        elif isinstance(data, PolarsDataFrame):
+            polars_df = data
+            columns = list(polars_df.columns)
+            rows = polars_df.to_numpy().tolist()
+            return [columns] + rows
+
+    if USE_PANDAS and isinstance(data, PandasDataFrame):  # type: ignore
+        # Convert pandas DataFrame to list of lists
+        pandas_df = data
+        columns = pandas_df.columns.tolist()
+        rows = pandas_df.values.tolist()
+        return [columns] + rows
+
+    # Assume it's a list of lists
+    return cast(list[list[str]], data)
diff --git a/src/pptxr/_pptx/slide.py b/src/pptxr/_pptx/slide.py
@@ -12,7 +12,7 @@
 from .converter import PptxConvertible, to_pptx_length
 from .shape import Shape
 from .shape.picture import Picture, PictureData, PictureProps
-from .shape.table import DataFrame, Table, TableData, TableProps
+from .shape.table import DataFrame, Table, TableData, TableProps, dataframe2list
 from .shape.text import Text, TextData, TextProps
 from .shape.title import Title
 
@@ -102,6 +102,7 @@ def picture(
         return self
 
     def table(self, data: DataFrame, **kwargs: Unpack[TableProps]) -> Self:
+        data = dataframe2list(data)
         rows, cols = len(data), len(data[0])
         table_data: TableData = {"type": "table", "data": data, **kwargs}
 
diff --git a/tests/test_table.py b/tests/test_table.py
@@ -0,0 +1,100 @@
+"""Tests for table module."""
+
+import pathlib
+
+import pytest
+
+import pptxr
+from pptxr._features import USE_PANDAS, USE_POLARS
+
+
+def test_create_table_with_list_data(output: pathlib.Path) -> None:
+    """Test creating a table with list data."""
+    table_data = [
+        ["Header 1", "Header 2", "Header 3"],
+        ["Cell 1,1", "Cell 1,2", "Cell 1,3"],
+        ["Cell 2,1", "Cell 2,2", "Cell 2,3"],
+    ]
+
+    presentation = (
+        pptxr.Presentation.builder()
+        .slide(
+            pptxr.SlideBuilder().table(
+                table_data,
+                left=(100, "pt"),
+                top=(100, "pt"),
+                width=(400, "pt"),
+                height=(200, "pt"),
+            )
+        )
+        .build()
+    )
+    presentation.save(output / "table_list_data.pptx")
+
+
+@pytest.mark.skipif(not USE_PANDAS, reason="Pandas not installed")
+def test_create_table_with_pandas_dataframe(output: pathlib.Path) -> None:
+    """Test creating a table with pandas DataFrame."""
+    # Type checking is done at runtime when pandas is available
+    import pandas as pd  # type: ignore[import]
+
+    # Create pandas DataFrame with simple data using dictionary
+    data = {
+        "名前": ["田中", "佐藤", "鈴木"],
+        "年齢": [25, 30, 22],
+        "都市": ["東京", "大阪", "名古屋"],
+    }
+    df = pd.DataFrame(data)
+
+    # Convert DataFrame to a list of lists for table creation
+    table_data = [df.columns.tolist()] + df.values.tolist()
+
+    presentation = (
+        pptxr.Presentation.builder()
+        .slide(
+            pptxr.SlideBuilder().table(
+                table_data,  # 変換したリストを使用
+                left=(100, "pt"),
+                top=(100, "pt"),
+                width=(400, "pt"),
+                height=(200, "pt"),
+                first_row_header=True,
+            )
+        )
+        .build()
+    )
+    presentation.save(output / "table_pandas_data.pptx")
+
+
+@pytest.mark.skipif(not USE_POLARS, reason="Polars not installed")
+def test_create_table_with_polars_dataframe(output: pathlib.Path) -> None:
+    """Test creating a table with polars DataFrame."""
+    # Type checking is done at runtime when polars is available
+    import polars as pl  # type: ignore[import]
+
+    # Create polars DataFrame with simple data using dictionary
+    data = {
+        "製品": ["A製品", "B製品", "C製品"],
+        "価格": [1000, 2000, 3000],
+        "在庫": [50, 30, 10],
+    }
+    df = pl.DataFrame(data)
+
+    # Convert DataFrame to a list of lists for table creation
+    table_data = [df.columns] + df.to_numpy().tolist()
+
+    presentation = (
+        pptxr.Presentation.builder()
+        .slide(
+            pptxr.SlideBuilder().table(
+                table_data,  # 変換したリストを使用
+                left=(100, "pt"),
+                top=(100, "pt"),
+                width=(400, "pt"),
+                height=(200, "pt"),
+                first_row_header=True,
+            )
+        )
+        .build()
+    )
+    presentation.save(output / "table_polars_data.pptx")
