Merge pull request #49 from daizutabi/feature

Add facet method to Plot class with flexible index and column handling

Author: daizutabi

src/xlviews/chart/axes.py

@@ -135,6 +135,14 @@ def copy(
         has_legend = self.chart.api[1].HasLegend
         include_in_layout = self.chart.api[1].Legend.IncludeInLayout
 
+        if left == 0:
+            left = self.chart.left + self.chart.width
+            top = self.chart.top
+
+        if top == 0:
+            left = self.chart.left
+            top = self.chart.top + self.chart.height
+
         return self.__class__(
             row=row,
             column=column,

src/xlviews/core/formula.py

@@ -70,14 +70,14 @@ def _aggregate(
     if func is None:
         return column
 
-    if isinstance(func, str) and func in AGG_FUNCS:
-        return f"AGGREGATE({AGG_FUNCS[func]},{option},{column})"
+    if isinstance(func, str):
+        if func in AGG_FUNCS:
+            return f"AGGREGATE({AGG_FUNCS[func]},{option},{column})"
 
-    if isinstance(func, Range | RangeImpl):
-        ref = func.get_address(column_absolute=False, row_absolute=False)
-    else:
-        ref = func
+        msg = f"Invalid aggregate function: {func}"
+        raise ValueError(msg)
 
+    ref = func.get_address(column_absolute=False, row_absolute=False)
     func = f"LOOKUP({ref},{{{AGG_FUNC_NAMES}}},{{{AGG_FUNC_INTS}}})"
     soa = aggregate("soa", ranges, option=option, **kwargs)
     return f'IF({ref}="soa",{soa},AGGREGATE({func},{option},{column}))'

src/xlviews/dataframes/heat_frame.py

@@ -198,11 +198,9 @@ def xs(
     columns: dict[Hashable, Any] | None,
 ) -> DataFrame:
     if index:
-        for key, value in index.items():
-            df = df.xs(value, level=key, axis=0)  # type: ignore
+        df = df.xs(tuple(index.values()), 0, tuple(index.keys()))  # type: ignore
 
     if columns:
-        for key, value in columns.items():
-            df = df.xs(value, level=key, axis=1)  # type: ignore
+        df = df.xs(tuple(columns.values()), 1, tuple(columns.keys()))  # type: ignore
 
     return df

src/xlviews/dataframes/style.py

@@ -264,11 +264,8 @@ def _set_heat_border(sf: HeatFrame) -> None:
     ec = rcParams["heat.border.color"]
 
     for row in iter_group_locs(sf.index, offset=r):
-        if row[0] == row[1]:
-            continue
-
         for col in iter_group_locs(sf.columns, offset=c):
-            if col[0] == col[1]:
+            if row[0] == row[1] and col[0] == col[1]:
                 continue
 
             rng = sf.sheet.range((row[0], col[0]), (row[1], col[1]))

src/xlviews/figure/groupby.py

@@ -5,6 +5,8 @@
 from itertools import cycle, islice
 from typing import TYPE_CHECKING, Generic, TypeAlias, TypeVar
 
+from pandas import MultiIndex
+
 from xlviews.chart.style import COLORS, MARKER_DICT
 
 if TYPE_CHECKING:
@@ -164,6 +166,9 @@ def get_palette(
     if style is None:
         return None
 
+    if data.index.name is not None or isinstance(data.index, MultiIndex):
+        data = data.index.to_frame(index=False)
+
     if isinstance(style, dict):
         return cls(data, data.columns.to_list(), style)
 
@@ -182,3 +187,17 @@ def get_palette(
         return cls(data, data.columns.tolist(), default)  # type: ignore
 
     return cls(data, columns)
+
+
+def get_marker_palette(
+    data: DataFrame,
+    marker: PaletteStyle | None,
+) -> MarkerPalette | None:
+    return get_palette(MarkerPalette, data, marker)  # type: ignore
+
+
+def get_color_palette(
+    data: DataFrame,
+    color: PaletteStyle | None,
+) -> ColorPalette | None:
+    return get_palette(ColorPalette, data, color)  # type: ignore

src/xlviews/figure/palette.py

@@ -5,13 +5,13 @@
 from typing import TYPE_CHECKING, TypeAlias
 
 import pandas as pd
-from pandas import DataFrame
+from pandas import DataFrame, Index
 
-from .palette import ColorPalette, MarkerPalette, PaletteStyle, get_palette
+from .palette import PaletteStyle, get_color_palette, get_marker_palette
 
 if TYPE_CHECKING:
-    from collections.abc import Iterator
-    from typing import Self
+    from collections.abc import Iterator, Sequence
+    from typing import Any, Self
 
     from xlviews.chart.axes import Axes
     from xlviews.chart.series import Series
@@ -70,8 +70,8 @@ def set(
         size: int | None = None,
     ) -> Self:
         index = self.data.index.to_frame(index=False)
-        marker_palette = get_palette(MarkerPalette, index, marker)
-        color_palette = get_palette(ColorPalette, index, color)
+        marker_palette = get_marker_palette(index, marker)
+        color_palette = get_color_palette(index, color)
 
         for key, s in zip(self.keys(), self.series_collection, strict=True):
             s.set(
@@ -85,6 +85,34 @@ def set(
 
         return self
 
+    @classmethod
+    def facet(
+        cls,
+        axes: Axes,
+        data: DataFrame,
+        index: str | list[str] | None = None,
+        columns: str | list[str] | None = None,
+    ) -> Iterator[tuple[dict[str, Any], Self]]:
+        left = axes.chart.left
+        top = axes.chart.top
+        width = axes.chart.width
+        height = axes.chart.height
+
+        for r, rkey in enumerate(iterrows(data.index, index)):
+            for c, ckey in enumerate(iterrows(data.index, columns)):
+                key = rkey | ckey
+                sub = xs(data, key)
+
+                if len(sub) == 0:
+                    continue
+
+                if r == 0 and c == 0:
+                    axes_ = axes
+                else:
+                    axes_ = axes.copy(left=left + c * width, top=top + r * height)
+
+                yield key, cls(axes_, sub)
+
 
 def get_label(label: Label, key: dict[str, Hashable]) -> str:
     if isinstance(label, str):
@@ -95,3 +123,29 @@ def get_label(label: Label, key: dict[str, Hashable]) -> str:
 
     msg = f"Invalid label: {label}"
     raise ValueError(msg)
+
+
+def iterrows(
+    index: Index,
+    levels: int | str | Sequence[int | str] | None,
+) -> Iterator[dict[str, Any]]:
+    if levels is None:
+        yield {}
+        return
+
+    if isinstance(levels, int | str):
+        levels = [levels]
+
+    if levels:
+        values = {level: index.get_level_values(level) for level in levels}
+        it = DataFrame(values).drop_duplicates().iterrows()
+
+        for _, s in it:
+            yield s.to_dict()
+
+
+def xs(df: DataFrame, index: dict[str, Any] | None) -> DataFrame:
+    if index:
+        df = df.xs(tuple(index.values()), 0, tuple(index.keys()), drop_level=False)  # type: ignore
+
+    return df

src/xlviews/figure/plot.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from xlviews.chart.axes import Axes
+from xlviews.figure.palette import get_color_palette, get_marker_palette
+from xlviews.figure.plot import Plot
+from xlviews.testing.common import create_sheet
+from xlviews.testing.sheet_frame.pivot import Pivot
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+    from typing import Any
+
+    from xlviews.dataframes.sheet_frame import SheetFrame
+
+
+def facet(sf: SheetFrame) -> Iterator[tuple[dict[str, Any], Plot]]:
+    sf.set_adjacent_column_width(1)
+
+    ax = Axes(2, 11)
+    data = sf.groupby(["A", "B", "X"]).agg(include_sheetname=True)
+    cp = get_color_palette(data, "X")
+    mp = get_marker_palette(data, "A")
+
+    for key, plot in Plot.facet(ax, data, index="B", columns="A"):
+        plot.add("u", "v").set(color=cp, marker=mp, label="{X}", alpha=0.8)
+        plot.axes.title = "{A}_{B}".format(**key)
+        yield key, plot
+
+
+if __name__ == "__main__":
+    sheet = create_sheet()
+    fc = Pivot(sheet, style=True)
+    list(facet(fc.sf))

src/xlviews/testing/figure/facet.py

@@ -9,12 +9,17 @@
 
 if TYPE_CHECKING:
     from collections.abc import Hashable, Iterator
-    from typing import Any
+    from typing import Any, Literal
 
     from xlviews.dataframes.sheet_frame import SheetFrame
 
 
-def pair(sf: SheetFrame) -> Iterator[tuple[dict[Hashable, Any], HeatFrame]]:
+def pair(
+    sf: SheetFrame,
+    values: str | list[str] | None = None,
+    columns: str | list[str] | None = None,
+    axis: Literal[0, 1] | None = None,
+) -> Iterator[tuple[dict[Hashable, Any], HeatFrame]]:
     sf.set_adjacent_column_width(1)
 
     rng = sf.get_range("u")
@@ -27,7 +32,15 @@ def pair(sf: SheetFrame) -> Iterator[tuple[dict[Hashable, Any], HeatFrame]]:
 
     df = sf.pivot_table(["u", "v"], ["B", "Y", "y"], ["A", "X", "x"], formula=True)
 
-    for key, frame in HeatFrame.pair(2, 13, df, index="B", columns="A", axis=1):
+    for key, frame in HeatFrame.pair(
+        2,
+        13,
+        df,
+        values=values,
+        index="B",
+        columns=columns,
+        axis=axis,
+    ):
         frame.autofit()
         frame.set_adjacent_column_width(1)
         if key["value"] == "u":
@@ -42,4 +55,6 @@ def pair(sf: SheetFrame) -> Iterator[tuple[dict[Hashable, Any], HeatFrame]]:
 if __name__ == "__main__":
     sheet = create_sheet()
     fc = Pivot(sheet, style=True)
-    list(pair(fc.sf))
+    # list(pair(fc.sf, values="u", columns=None))
+    # list(pair(fc.sf, values=None, columns="A", axis=1))
+    list(pair(fc.sf, values="v", columns="A"))

src/xlviews/testing/heat_frame/pair.py

@@ -67,3 +67,15 @@ def test_chart_position_top_left(sheet: Sheet):
     c = Axes(sheet=sheet, top=0, left=200)
     assert c.chart.left == 200
     assert c.chart.top == TOP
+
+
+def test_copy_right(sheet: Sheet):
+    a = Axes(sheet=sheet)
+    b = a.copy(left=0)
+    assert b.chart.left == a.chart.left + a.chart.width
+
+
+def test_copy_bottom(sheet: Sheet):
+    a = Axes(sheet=sheet)
+    b = a.copy(top=0)
+    assert b.chart.top == a.chart.top + a.chart.height