plot-io 0.1.0__py3-none-any.whl

plot_io/decorations.py

@@ -0,0 +1,496 @@
+"""Decorations: reference lines, spans, text and annotations on a plot.
+
+Decorations are not series: they take no palette color, are in the legend only with a
+``name``, and change the axis limits only with ``extend_limits=True``.  They are drawn
+with the series in ``z`` order (spans under the series, lines and text over them by
+default).  Positions are in data units (``coords="data"``), fractions of the axes
+(``"axes"``) or of the figure (``"figure"``).
+"""
+
+import math
+from typing import Any
+
+from .settings import (Setting, Settings, boolean, color, number, one_of, table, text)
+from .style import dash_of, resolve_color, with_alpha
+from .text import Font, draw_text, measure
+
+__all__ = ["Decoration", "HLine", "VLine", "AxLine", "HSpan", "VSpan", "Text", "Annotate"]
+
+
+def _dash(v: Any) -> Any:
+    if isinstance(v, (list, tuple)):
+        return [number(lo=0)(x) for x in v]
+    return one_of("solid", "dashed", "dotted", "dashdot")(v)
+
+
+def _point(v: Any) -> tuple[Any, Any]:
+    if not isinstance(v, (tuple, list)) or len(v) != 2:
+        raise ValueError(f"expected a point (x, y), got {v!r}")
+    return v[0], v[1]
+
+
+_COORDS = one_of("data", "axes", "figure")
+
+
+def _arrow_style(v: Any) -> str:
+    if v not in ("->", "-|>", "-"):
+        raise ValueError(f"expected '->', '-|>' or '-', got {v!r}")
+    return v
+
+
+def _common() -> dict[str, Setting]:
+    return dict(
+        name=Setting(text, doc="Legend entry; unset: not in the legend."),
+        alpha=Setting(number(0, 1), doc="Opacity."),
+        z=Setting(number(), doc="Drawing order among the series (default 0)."),
+        extend_limits=Setting(boolean, doc="Include the decoration in the axis limits."),
+        clip=Setting(boolean, doc="Clip to the axes area."),
+    )
+
+
+def _line() -> dict[str, Setting]:
+    return dict(
+        color=Setting(color, doc="Line color."),
+        line_width=Setting(number(lo=0), inherited=True, doc="Line width (pt).", length=True),
+        line_style=Setting(_dash,
+                           doc="`solid`, `dashed`, `dotted`, `dashdot` or a dash list."),
+    )
+
+
+def _font() -> dict[str, Setting]:
+    return dict(
+        font_size=Setting(number(positive=True), inherited=True, doc="Font size (pt).",
+                          length=True),
+        font_family=Setting(str, inherited=True, doc="Font family."),
+        font_weight=Setting(one_of("normal", "bold"), doc="`normal` or `bold`."),
+        text_color=Setting(color, inherited=True, doc="Text color."),
+    )
+
+
+class Decoration(Settings):
+    """Base of decorations: behaves like a series towards the axes and the legend."""
+
+    USES_PALETTE = False
+    AXES = True
+    DODGE_GROUP = None
+    STACK_GROUP = None
+    #: The axis roles the decoration's data coordinates are on.
+    ROLES: tuple[str, ...] = ()
+
+    def __init__(self, plot: Any, **settings: Any) -> None:
+        self.plot = plot
+        super().__init__(**settings)
+        plot.decorations.append(self)
+
+    def _parent(self) -> Settings:
+        return self.plot
+
+    def __repr__(self) -> str:
+        return f"{self.kind}()"
+
+    # -- what the axes need (only with extend_limits) -------------------------------
+
+    def has_axis(self, role: str) -> bool:
+        return role in self.ROLES and self.get("extend_limits")
+
+    def categories(self, role: str) -> list[Any] | None:
+        return None
+
+    def data_values(self, role: str) -> list[float]:
+        return []
+
+    def values(self, role: str) -> Any:
+        from array import array
+        return array("d", [float(v) for v in self.data_values(role)
+                           if isinstance(v, (int, float))])
+
+    def extent(self, role: str, scale: Any) -> tuple[int, float, float]:
+        from . import _backend
+        return _backend.extent(self.values(role), scale.function)
+
+    def inverts(self, role: str) -> bool:
+        return False
+
+    def zero_default(self, role: str) -> bool:
+        return False
+
+    def zero_padding(self, role: str) -> bool:
+        return False
+
+    def color_sources(self) -> list[Any]:
+        return []
+
+    def size_sources(self) -> list[Any]:
+        return []
+
+    def has_own_color(self) -> bool:
+        return True
+
+    def legend_labels(self) -> list[Any]:
+        """Legend entries of the object: its ``name``, or none."""
+        name = self.get("name")
+        return [] if name is None else [name]
+
+    def legend_glyph_at(self, ctx: Any, i: int, x: float, y: float, w: float, h: float) -> None:
+        """The symbol of legend entry ``i`` (series with several entries override this)."""
+        self.legend_glyph(ctx, x, y, w, h)
+
+    def legend_colors_at(self, ctx: Any, i: int) -> tuple[Any, Any]:
+        return self.legend_colors(ctx)
+
+    def legend_points(self) -> Any:
+        return None
+
+    def legend_colors(self, ctx: Any) -> tuple[Any, Any]:
+        c = with_alpha(ctx.color(self.get("color")), self.get("alpha"))
+        return c, None
+
+    def legend_glyph(self, ctx: Any, x: float, y: float, w: float, h: float) -> None:
+        c = with_alpha(ctx.color(self.get("color")), self.get("alpha"))
+        if c is None:
+            return
+        if "line_style" in self.SETTINGS:
+            lw = self.get("line_width")
+            ctx.canvas.draw_lines([x, x + w], [y + h / 2, y + h / 2], color=c, width=lw,
+                                  dash=dash_of(self.get("line_style"), lw))
+        else:
+            ctx.canvas.fill_style = c
+            ctx.canvas.fill_rect(x, y, w, h)
+
+    # -- positions ------------------------------------------------------------------
+
+    def _pos(self, ctx: Any, role: str, v: Any, coords: str) -> float:
+        """Page coordinate of ``v`` along ``role`` (``x`` or ``y``)."""
+        x0, y0, x1, y1 = ctx.rect
+        if coords == "axes":
+            return x0 + float(v) * (x1 - x0) if role == "x" else y1 - float(v) * (y1 - y0)
+        if coords == "figure":
+            fw, fh = getattr(ctx, "figure_size", (x1, y1))
+            return float(v) * fw if role == "x" else (1.0 - float(v)) * fh
+        state = ctx.x if role == "x" else ctx.y
+        if isinstance(v, str):
+            if not state.is_category or v not in state.categories:
+                raise ValueError(f"{v!r} is not a category of the {role} axis")
+            v = state.categories.index(v)
+        mapper = ctx.xmap if role == "x" else ctx.ymap
+        return mapper(float(v))
+
+    def draw(self, ctx: Any) -> None:
+        raise NotImplementedError
+
+
+class _Rule(Decoration):
+    """A horizontal or vertical reference line across part of the axes."""
+
+    VERTICAL = False
+    value: Any = None
+
+    def draw(self, ctx: Any) -> None:
+        c = with_alpha(ctx.color(self.get("color")), self.get("alpha"))
+        lw = self.get("line_width")
+        if c is None:
+            return
+        x0, y0, x1, y1 = ctx.rect
+        lo, hi = self.get("min"), self.get("max")
+        if self.VERTICAL:
+            p = self._pos(ctx, "x", self.value, "data")
+            xs, ys = [p, p], [y1 - lo * (y1 - y0), y1 - hi * (y1 - y0)]
+        else:
+            p = self._pos(ctx, "y", self.value, "data")
+            xs, ys = [x0 + lo * (x1 - x0), x0 + hi * (x1 - x0)], [p, p]
+        if not all(math.isfinite(v) for v in xs + ys):
+            return
+        ctx.canvas.draw_lines(xs, ys, color=c, width=lw, dash=dash_of(self.get("line_style"), lw),
+                              cap="butt")
+
+    def data_values(self, role: str) -> list[float]:
+        return [self.value]
+
+
+class HLine(_Rule):
+    """A horizontal line at ``y`` from ``min`` to ``max`` (fractions of the axes width)."""
+
+    SETTINGS = table(
+        **_line(),
+        min=Setting(number(0, 1), doc="Start, as a fraction of the axes width."),
+        max=Setting(number(0, 1), doc="End, as a fraction of the axes width."),
+        **_common(),
+    )
+    ROLES = ("y",)
+
+    def __init__(self, plot: Any, y: Any, **settings: Any) -> None:
+        self.value = y
+        super().__init__(plot, **settings)
+
+
+class VLine(_Rule):
+    """A vertical line at ``x`` from ``min`` to ``max`` (fractions of the axes height)."""
+
+    SETTINGS = HLine.SETTINGS
+    ROLES = ("x",)
+    VERTICAL = True
+
+    def __init__(self, plot: Any, x: Any, **settings: Any) -> None:
+        self.value = x
+        super().__init__(plot, **settings)
+
+
+class AxLine(Decoration):
+    """An infinite straight line through ``point`` and ``point2``, or with ``slope`` (data
+    units; linear axes)."""
+
+    SETTINGS = table(**_line(), **_common())
+    ROLES = ("x", "y")
+
+    def __init__(self, plot: Any, point: Any, slope: Any = None, point2: Any = None,
+                 **settings: Any) -> None:
+        self.p1 = _point(point)
+        if (slope is None) == (point2 is None):
+            raise ValueError("give either slope or point2")
+        self.p2 = _point(point2) if point2 is not None else None
+        self.slope = None if slope is None else float(slope)
+        super().__init__(plot, **settings)
+
+    def data_values(self, role: str) -> list[float]:
+        i = 0 if role == "x" else 1
+        pts = [self.p1] + ([self.p2] if self.p2 is not None else [])
+        return [p[i] for p in pts]
+
+    def draw(self, ctx: Any) -> None:
+        c = with_alpha(ctx.color(self.get("color")), self.get("alpha"))
+        lw = self.get("line_width")
+        if c is None:
+            return
+        if self.p2 is None:
+            for st in (ctx.x, ctx.y):
+                if st.scale is None or st.scale.name != "linear":
+                    raise ValueError("axline with a slope needs linear axes; give point2")
+            p2 = (float(self.p1[0]) + 1.0, float(self.p1[1]) + float(self.slope or 0.0))
+        else:
+            p2 = self.p2
+        ax, ay = self._pos(ctx, "x", self.p1[0], "data"), self._pos(ctx, "y", self.p1[1], "data")
+        bx, by = self._pos(ctx, "x", p2[0], "data"), self._pos(ctx, "y", p2[1], "data")
+        dx, dy = bx - ax, by - ay
+        n = math.hypot(dx, dy)
+        if not n or not math.isfinite(n):
+            return
+        x0, y0, x1, y1 = ctx.rect
+        far = 2.0 * (math.hypot(x1 - x0, y1 - y0) + abs(ax - x0) + abs(ay - y0)) / n
+        ctx.canvas.draw_lines([ax - far * dx, ax + far * dx], [ay - far * dy, ay + far * dy],
+                              color=c, width=lw, dash=dash_of(self.get("line_style"), lw))
+
+
+class _Span(Decoration):
+    VERTICAL = False
+    v0: Any = None
+    v1: Any = None
+
+    def draw(self, ctx: Any) -> None:
+        c = with_alpha(ctx.color(self.get("color")), self.get("alpha"))
+        if c is None:
+            return
+        x0, y0, x1, y1 = ctx.rect
+        lo, hi = self.get("min"), self.get("max")
+        if self.VERTICAL:
+            a, b = self._pos(ctx, "x", self.v0, "data"), self._pos(ctx, "x", self.v1, "data")
+            r = (min(a, b), y1 - hi * (y1 - y0), abs(b - a), (hi - lo) * (y1 - y0))
+        else:
+            a, b = self._pos(ctx, "y", self.v0, "data"), self._pos(ctx, "y", self.v1, "data")
+            r = (x0 + lo * (x1 - x0), min(a, b), (hi - lo) * (x1 - x0), abs(b - a))
+        if not all(math.isfinite(v) for v in r):
+            return
+        canvas = ctx.canvas
+        canvas.begin_path()
+        canvas.rect(*r)
+        canvas.fill(c)
+        lc = ctx.color(self.get("line_color"))
+        if lc is not None and self.get("line_width"):
+            canvas.line_width = self.get("line_width")
+            canvas.stroke(lc)
+
+    def data_values(self, role: str) -> list[float]:
+        return [self.v0, self.v1]
+
+
+class HSpan(_Span):
+    """A horizontal band from ``y0`` to ``y1``, from ``min`` to ``max`` of the axes width."""
+
+    SETTINGS = table(
+        color=Setting(color, doc="Fill color."),
+        line_color=Setting(color, doc="Outline color; default none."),
+        line_width=Setting(number(lo=0), doc="Outline width (pt).", length=True),
+        min=Setting(number(0, 1), doc="Start, as a fraction of the axes width."),
+        max=Setting(number(0, 1), doc="End, as a fraction of the axes width."),
+        **_common(),
+    )
+    ROLES = ("y",)
+
+    def __init__(self, plot: Any, y0: Any, y1: Any, **settings: Any) -> None:
+        self.v0, self.v1 = y0, y1
+        super().__init__(plot, **settings)
+
+
+class VSpan(_Span):
+    """A vertical band from ``x0`` to ``x1``, from ``min`` to ``max`` of the axes height."""
+
+    SETTINGS = HSpan.SETTINGS
+    ROLES = ("x",)
+    VERTICAL = True
+
+    def __init__(self, plot: Any, x0: Any, x1: Any, **settings: Any) -> None:
+        self.v0, self.v1 = x0, x1
+        super().__init__(plot, **settings)
+
+
+class Text(Decoration):
+    """Text at ``(x, y)``: a string or rich content."""
+
+    SETTINGS = table(
+        align=Setting(one_of("left", "center", "right"),
+                      doc="Horizontal alignment on the point."),
+        valign=Setting(one_of("baseline", "top", "center", "bottom"),
+                       doc="Vertical alignment on the point."),
+        rotation=Setting(number(), doc="Rotation (degrees, counterclockwise)."),
+        coords=Setting(_COORDS, doc="`data`, `axes` (0–1) or `figure` (0–1)."),
+        color=Setting(color, doc="Text color; default `text_color`."),
+        **_font(),
+        **{k: v for k, v in _common().items() if k != "clip"},
+        clip=Setting(boolean, doc="Clip to the axes area."),
+    )
+    ROLES = ("x", "y")
+
+    def __init__(self, plot: Any, x: Any, y: Any, content: Any, **settings: Any) -> None:
+        self.x, self.y, self.content = x, y, content
+        super().__init__(plot, **settings)
+
+    def has_axis(self, role: str) -> bool:
+        return super().has_axis(role) and self.get("coords") == "data"
+
+    def data_values(self, role: str) -> list[float]:
+        return [self.x if role == "x" else self.y]
+
+    def font(self) -> Font:
+        c = self.get("color")
+        return Font(self.get("font_size"), self.get("font_family"), self.get("font_weight"),
+                    color=c)
+
+    def legend_colors(self, ctx: Any) -> tuple[Any, Any]:
+        return ctx.color(self.get("color")), None
+
+    def draw(self, ctx: Any) -> None:
+        coords = self.get("coords")
+        px, py = self._pos(ctx, "x", self.x, coords), self._pos(ctx, "y", self.y, coords)
+        font = self.font()
+        font.color = with_alpha(resolve_color(font.color, ctx.palette), self.get("alpha"))
+        draw_text(ctx.canvas, self.content, px, py, font, self.get("align"), self.get("valign"),
+                  self.get("rotation"))
+
+
+class Annotate(Text):
+    """Text at ``xytext`` with an arrow to ``xy``."""
+
+    SETTINGS = table(
+        **{k: v for k, v in Text.SETTINGS.items() if k not in ("align", "valign")},
+        align=Setting(one_of("left", "center", "right"),
+                      doc="Horizontal alignment on `xytext`."),
+        valign=Setting(one_of("baseline", "top", "center", "bottom"),
+                       doc="Vertical alignment on `xytext`."),
+        text_coords=Setting(one_of("data", "axes", "figure", "offset"),
+                            doc="Coordinates of `xytext`: as `coords`, or `offset` (pt from "
+                                "`xy`, y up); default `coords`."),
+        arrow=Setting(boolean, doc="Draw an arrow from the text to `xy`."),
+        arrow_style=Setting(_arrow_style,
+                            doc="`->` (open head), `-|>` (filled head) or `-` (no head)."),
+        arrow_color=Setting(color, doc="Arrow color; default the text color."),
+        arrow_width=Setting(number(lo=0), doc="Arrow line width (pt).", length=True),
+        head_size=Setting(number(lo=0), doc="Arrowhead length (pt).", length=True),
+        shrink=Setting(number(lo=0), doc="Gap left at both ends of the arrow (pt).",
+                       length=True),
+    )
+
+    def __init__(self, plot: Any, content: Any, xy: Any, xytext: Any = None,
+                 **settings: Any) -> None:
+        self.xy = _point(xy)
+        self.xytext = _point(xytext) if xytext is not None else None
+        x, y = self.xytext if self.xytext is not None else self.xy
+        super().__init__(plot, x, y, content, **settings)
+
+    def data_values(self, role: str) -> list[float]:
+        i = 0 if role == "x" else 1
+        out = [self.xy[i]]
+        if self.xytext is not None and self.get("text_coords") == "data":
+            out.append(self.xytext[i])
+        return out
+
+    def draw(self, ctx: Any) -> None:
+        coords = self.get("coords")
+        tx = self.get("text_coords")
+        ax, ay = self._pos(ctx, "x", self.xy[0], coords), self._pos(ctx, "y", self.xy[1], coords)
+        if self.xytext is None:
+            px, py = ax, ay
+        elif tx == "offset":
+            px, py = ax + float(self.xytext[0]), ay - float(self.xytext[1])
+        else:
+            px, py = self._pos(ctx, "x", self.xytext[0], tx), self._pos(ctx, "y", self.xytext[1], tx)
+        font = self.font()
+        canvas = ctx.canvas
+        tc = with_alpha(resolve_color(font.color, ctx.palette), self.get("alpha"))
+        font.color = tc
+        w, a, d = measure(canvas, self.content, font)
+        bx = px + {"left": 0.0, "center": -w / 2, "right": -w}[self.get("align")]
+        top = py - {"baseline": a, "top": 0.0, "center": (a + d) / 2, "bottom": a + d}[
+            self.get("valign")]
+        draw_text(canvas, self.content, px, py, font, self.get("align"), self.get("valign"))
+        if self.get("arrow") and self.xytext is not None:
+            ac = with_alpha(ctx.color(self.get("arrow_color")), self.get("alpha"))
+            self._arrow(canvas, (bx, top, bx + w, top + a + d), ax, ay, ac)
+
+    def _arrow(self, canvas: Any, box: tuple[float, float, float, float], ax: float, ay: float,
+               c: Any) -> None:
+        x0, y0, x1, y1 = box
+        cx, cy = (x0 + x1) / 2, (y0 + y1) / 2
+        dx, dy = ax - cx, ay - cy
+        n = math.hypot(dx, dy)
+        if not n or c is None:
+            return
+        ux, uy = dx / n, dy / n
+        # leave the text box: smallest t where the ray crosses its edge
+        ts = []
+        if ux:
+            ts.append(((x1 if ux > 0 else x0) - cx) / ux)
+        if uy:
+            ts.append(((y1 if uy > 0 else y0) - cy) / uy)
+        t0 = min(t for t in ts if t >= 0) if ts else 0.0
+        shrink = self.get("shrink")
+        t0 += shrink
+        t1 = n - shrink
+        if t1 <= t0:
+            return
+        sx, sy = cx + ux * t0, cy + uy * t0
+        ex, ey = cx + ux * t1, cy + uy * t1
+        lw = self.get("arrow_width")
+        style = self.get("arrow_style")
+        hs = self.get("head_size")
+        if style == "-|>":
+            ex2, ey2 = ex - ux * hs * 0.9, ey - uy * hs * 0.9
+        else:
+            ex2, ey2 = ex, ey
+        canvas.draw_lines([sx, ex2], [sy, ey2], color=c, width=lw, cap="butt")
+        if style == "-" or not hs:
+            return
+        hw = 0.45 * hs
+        px_, py_ = -uy, ux
+        bx, by = ex - ux * hs, ey - uy * hs
+        canvas.save()
+        canvas.begin_path()
+        canvas.move_to(bx + px_ * hw, by + py_ * hw)
+        canvas.line_to(ex, ey)
+        canvas.line_to(bx - px_ * hw, by - py_ * hw)
+        if style == "-|>":
+            canvas.close_path()
+            canvas.fill(c)
+        else:
+            canvas.line_width = lw
+            canvas.line_join = "miter"
+            canvas.stroke(c)
+        canvas.restore()