Merge pull request #51 from daizutabi/50-color-and-marker-accept-callable

Enhance Plot and Palette modules with flexible marker and color generation

Author: daizutabi

src/xlviews/figure/palette.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from collections.abc import Hashable
+from collections.abc import Callable, Hashable
 from itertools import cycle, islice
 from typing import TYPE_CHECKING, Generic, TypeAlias, TypeVar
 
@@ -102,12 +102,11 @@ def get(self, value: Hashable) -> int:
 
         return self.index[value]
 
-    def __getitem__(self, value: Hashable | dict) -> T:
-        if value == {None: 0}:  # from series
+    def __getitem__(self, key: dict) -> T:
+        if key == {None: 0}:  # from series
             return self.items[0]
 
-        if isinstance(value, dict):
-            value = tuple(value[k] for k in self.columns)
+        value = tuple(key[k] for k in self.columns)
 
         return self.items[self.get(value)]
 
@@ -145,35 +144,63 @@ def cycle_colors(skips: Iterable[str] | None = None) -> Iterator[str]:
             yield color
 
 
+class FunctionPalette(Generic[T]):
+    columns: str | list[str]
+    func: Callable[[Hashable], T]
+
+    def __init__(self, columns: str | list[str], func: Callable[[Hashable], T]) -> None:
+        self.columns = columns
+        self.func = func
+
+    def __getitem__(self, key: dict) -> T:
+        if isinstance(self.columns, str):
+            return self.func(key[self.columns])
+
+        value = tuple(key[k] for k in self.columns)
+        return self.func(value)
+
+
 PaletteStyle: TypeAlias = (
     str
     | list[str]
     | dict[Hashable, str]
-    | tuple[str | list[str], dict[Hashable, str] | list[str]]
+    | Callable[[Hashable], str]
+    | tuple[str | list[str], list[str] | dict[Hashable, str]]
+    | tuple[str | list[str], Callable[[Hashable], str]]
     | Palette
+    | FunctionPalette
 )
 
 
 def get_palette(
     cls: type[Palette],
     data: DataFrame,
     style: PaletteStyle | None,
-) -> Palette | None:
+) -> Palette | FunctionPalette | None:
     """Get a palette from a style."""
-    if isinstance(style, Palette):
+    if isinstance(style, Palette | FunctionPalette):
         return style
 
     if style is None:
         return None
 
+    if isinstance(style, Callable):
+        if isinstance(data.index, MultiIndex):
+            return FunctionPalette(data.index.names, style)
+        return FunctionPalette(data.index.name, style)
+
     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)
 
     if isinstance(style, tuple):
-        return cls(data, *style)
+        columns, default = style
+        if callable(default):
+            return FunctionPalette(columns, default)
+
+        return cls(data, columns, default)
 
     columns = style
 
@@ -192,12 +219,12 @@ def get_palette(
 def get_marker_palette(
     data: DataFrame,
     marker: PaletteStyle | None,
-) -> MarkerPalette | None:
+) -> MarkerPalette | FunctionPalette | None:
     return get_palette(MarkerPalette, data, marker)  # type: ignore
 
 
 def get_color_palette(
     data: DataFrame,
     color: PaletteStyle | None,
-) -> ColorPalette | None:
+) -> ColorPalette | FunctionPalette | None:
     return get_palette(ColorPalette, data, color)  # type: ignore

src/xlviews/figure/plot.py

@@ -69,9 +69,8 @@ def set(
         weight: float | None = None,
         size: int | None = None,
     ) -> Self:
-        index = self.data.index.to_frame(index=False)
-        marker_palette = get_marker_palette(index, marker)
-        color_palette = get_color_palette(index, color)
+        marker_palette = get_marker_palette(self.data, marker)
+        color_palette = get_color_palette(self.data, color)
 
         for key, s in zip(self.keys(), self.series_collection, strict=True):
             s.set(

src/xlviews/testing/figure/plot.py

@@ -1,12 +1,17 @@
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 from xlwings.constants import ChartType
 
 from xlviews.chart.axes import Axes
 from xlviews.figure.plot import Plot
 from xlviews.testing.chart import Base
 from xlviews.testing.common import create_sheet
 
+if TYPE_CHECKING:
+    from collections.abc import Hashable
+
 if __name__ == "__main__":
     sheet = create_sheet()
     fc = Base(sheet, style=True)
@@ -21,23 +26,49 @@
         .set(label="abc", marker="o", color="blue", alpha=0.6)
     )
 
-    ax = Axes(chart_type=ChartType.xlXYScatterLines)
+    ax = Axes()
     data = sf.groupby("b").agg(include_sheetname=True)
     p = (
         Plot(ax, data)
-        .add("x", "y")
+        .add("x", "y", ChartType.xlXYScatterLines)
         .set(label="b={b}", marker=["o", "s"], color={"s": "red", "t": "blue"})
     )
 
-    ax = Axes(chart_type=ChartType.xlXYScatterLines)
+    ax = Axes()
     data = sf.groupby(["b", "c"]).agg(include_sheetname=True)
     p = (
         Plot(ax, data)
-        .add("x", "y")
+        .add("x", "y", ChartType.xlXYScatterLines)
         .set(
             label=lambda x: f"{x['b']},{x['c']}",
             marker="b",
             color=("c", ["red", "green"]),
             size=10,
         )
     )
+
+    def m(x: Hashable) -> str:
+        if x == "s":
+            return "o"
+        return "^"
+
+    ax = Axes(left=0)
+    data = sf.groupby("b").agg(include_sheetname=True)
+    p = (
+        Plot(ax, data)
+        .add("x", "y", ChartType.xlXYScatter)
+        .set(label="{b}", marker=m, size=10)
+    )
+
+    def c(x: Hashable) -> str:
+        if x == ("s", 100):
+            return "red"
+        return "blue"
+
+    ax = Axes()
+    data = sf.groupby(["b", "c"]).agg(include_sheetname=True)
+    p = (
+        Plot(ax, data)
+        .add("x", "y", ChartType.xlXYScatter)
+        .set(label="{b}_{c}", color=c, marker=("b", m), size=10)
+    )

tests/dataframes/dist_frame/test_sigma.py

@@ -4,6 +4,9 @@
 from xlwings import Sheet
 
 from xlviews.core.range import Range
+from xlviews.testing import is_app_available
+
+pytestmark = pytest.mark.skipif(not is_app_available(), reason="Excel not installed")
 
 
 @pytest.fixture(params=["norm", "weibull"])

tests/figure/test_grid.py

@@ -3,6 +3,9 @@
 
 from xlviews.chart.axes import Axes
 from xlviews.figure.grid import Grid, Series
+from xlviews.testing import is_app_available
+
+pytestmark = pytest.mark.skipif(not is_app_available(), reason="Excel not installed")
 
 
 @pytest.fixture(scope="module")