From d95fdb9124f5e3bc69eab0df90c5e5ce995cf7ef Mon Sep 17 00:00:00 2001 From: testycool Date: Sat, 10 Oct 2026 02:40:32 +0300 Subject: [PATCH] Let other packages supply a mark look without changing AInotate People want marks that look different from AInotate's bold default, for example a calm hand-drawn style. Building each style into AInotate would add fonts and drawing code to every install. A plug-in point keeps AInotate small: a look lives in its own package and is used only where installed. render() now draws every mark through a Look object. The built-in look, the only one shipped, is `ainotate.looks.Look` (also importable from ainotate.render). It holds exactly the drawing calls render() made before: colors, label font and size, labels and text notes, outlines, step badges, arrows, highlights, the spotlight dim and corner, the redaction fill, the keycap font, the magnifier line width and the drop shadow. A package registers a subclass under the entry point group "ainotate.looks" and overrides what it changes. A look can change a redaction's color, never its area. The look is chosen by spec "look" (also --look on annotate and shoot, and the `look` argument of the MCP annotate, preview and shoot tools), then the env AINOTATE_LOOK, then `look` in config.toml, then "default", the same order as the output folder. A look never breaks a render: when it is not installed, cannot be imported, raises while being constructed or raises while drawing, the default look is used and the render gets a warning that names the look and the error, or lists the installed looks. The warning reaches users like every render warning: result.warnings, the CLI's WARN lines and the MCP result's warnings. A failure while drawing redraws the whole image with the default look, so a half-drawn mix never ships. A look can never change a redaction or the caller's spec. An installed look gets its own copies of the page, marks, rects and points, and its own copy of the page image in every hook that is handed it, so nothing it changes reaches the render, the redo or the spec. Every value render() reads from it (attributes such as shadow and stroke, and the results of color, redact_fill, label_size and keys_font) is read and checked inside the same guard, so an attribute that raises or a value Pillow cannot draw ("not-a-color"), or a number that is not finite, is a failure of the look, with the same fallback; a color name such as "info" is passed on as the color it names. As a last line, when an installed look is in use and anything else raises, even AInotate's own code on a value the look gave it, the image is redrawn from the spec with the default look and a warning. An error the default look hits as well is AInotate's or the spec's, and is raised as before. The spec reference, the renderer notes served as ainotate://spec, the MCP fallback summary and a description on the MCP `look` parameter say how to pick a look. Testing: the default look is pixel-identical. A test pins the sha256 of a render of every mark type, light and dark, from before this change (Pillow 12.3.0 / FreeType 2.14.3; other versions skip the pin). The golden label-position test is unchanged but needs macOS Helvetica and skips on Linux; run with DejaVu Sans Bold instead, its three specs give the same label boxes and identical images on origin/main and this branch. New tests register a fake look and check that it draws, that the flag, env and config pick it, each fallback with its exact warning, the CLI and MCP output, the MCP schema text, that looks which change their geometry, spec data or page image leave redactions and the spec intact, that a raising attribute, bad hook results, an infinite shadow and a label size no frame can hold fall back, that a color name works, and that an error the default look also hits is still raised. Co-Authored-By: Claude Opus 5.5 --- skills/ainotate/references/spec.md | 15 +- src/ainotate/cli.py | 4 + src/ainotate/looks.py | 191 +++++++++++++++++++ src/ainotate/marks.py | 10 +- src/ainotate/mcp_docs.py | 3 +- src/ainotate/mcp_server.py | 31 ++- src/ainotate/output.py | 8 + src/ainotate/render.py | 84 +++++---- src/ainotate/spec.py | 2 + tests/conftest.py | 3 +- tests/test_looks.py | 294 +++++++++++++++++++++++++++++ 11 files changed, 594 insertions(+), 51 deletions(-) create mode 100644 src/ainotate/looks.py create mode 100644 tests/test_looks.py diff --git a/skills/ainotate/references/spec.md b/skills/ainotate/references/spec.md index 4aaf10c..1c22e43 100644 --- a/skills/ainotate/references/spec.md +++ b/skills/ainotate/references/spec.md @@ -16,6 +16,7 @@ Everything the renderer accepts. `SKILL.md` has the workflow; this file is for l | `crop_pad` | margin around the marks for `"auto"` and `"tight"` (default 260 image px at 2000px width) | | `dim` | spotlight darkness 0-1 (default 0.55) | | `arrow_style` | `"skitch"` (default, tapered), `"straight"`, `"curved"`, or `"line"` | +| `look` | mark style: `"default"` (built in) or an installed look package such as `"neat"`; see below | | `frame` | backdrop and window chrome, see below | | `privacy` | `"off"` (default for images), `"auto"` (OCR the image, redact findings) or an object | @@ -32,6 +33,14 @@ Everything the renderer accepts. `SKILL.md` has the workflow; this file is for l | `browser` | `chromium` (default), `firefox`, `webkit` (launch mode only) | | `timeout_ms` | navigation and wait timeout (default 30000) | +**Looks.** `"look"` picks the mark style. `"default"` is the only one built in; other packages add +looks (entry point group `ainotate.looks`, a subclass of `ainotate.looks.Look`), for example the +separate `ainotate-neat` package adds `"neat"`. Pick one per spec, with `--look` on `annotate` and +`shoot`, or with the MCP `look` argument; without one, env `AINOTATE_LOOK`, then `look` in +`~/.config/ainotate/config.toml`, then `"default"`. A look that is not installed, or fails to +load, start or draw, never breaks a render: the default look is used and a warning says why and +lists the installed looks. + `shoot` defaults `crop` to `"auto"` and `privacy` to `"auto"`. With chrome `"browser"` and no `url` in the frame, the address bar shows the captured page's address. @@ -149,9 +158,9 @@ explains it below, and can use `**bold**`, `` `code` `` and lines starting with Final images: `/ .png`, plus an identical copy in `backup_dir` if set. Never overwrites (adds " (2)"). Configuration, highest first: -env `AINOTATE_OUTPUT_DIR`, `AINOTATE_BACKUP_DIR`, `AINOTATE_PREFIX`; then -`~/.config/ainotate/config.toml` keys `output_dir`, `backup_dir`, `prefix`; then defaults -`~/Pictures/AInotate`, no backup, prefix `AInotate`. +env `AINOTATE_OUTPUT_DIR`, `AINOTATE_BACKUP_DIR`, `AINOTATE_PREFIX`, `AINOTATE_LOOK`; then +`~/.config/ainotate/config.toml` keys `output_dir`, `backup_dir`, `prefix`, `look`; then defaults +`~/Pictures/AInotate`, no backup, prefix `AInotate`, look `default`. A spec's `look` beats both. ## More patterns diff --git a/src/ainotate/cli.py b/src/ainotate/cli.py index 11779bc..5491edd 100644 --- a/src/ainotate/cli.py +++ b/src/ainotate/cli.py @@ -138,6 +138,8 @@ def _result_args(privacy_default): return [A("spec", help="spec.json, inline JSON, or - for stdin"), A("--draft", "--no-save", dest="draft", action="store_true", help="write to a temp dir"), A("--privacy", choices=("auto", "off"), help=f"solid redaction (default: {privacy_default})"), + A("--look", help="mark style: default or an installed look such as neat " + "(default: AINOTATE_LOOK, then look in config.toml)"), A("--copy", action="store_true", help="also put the image on the clipboard"), A("--json", action="store_true", help="print JSON status on stdout"), A("--debug", action="store_true", help="draw target/label boxes; implies --draft")] @@ -147,6 +149,8 @@ def _spec_with_flags(a): spec = _load_spec(a.spec) if a.privacy and isinstance(spec, dict): spec["privacy"] = a.privacy + if a.look and isinstance(spec, dict): + spec["look"] = a.look return spec diff --git a/src/ainotate/looks.py b/src/ainotate/looks.py new file mode 100644 index 0000000..9cfae69 --- /dev/null +++ b/src/ainotate/looks.py @@ -0,0 +1,191 @@ +"""Mark looks. `Look` is the built-in look and the base class for looks from other packages: +entry point group "ainotate.looks", name -> a Look subclass. The look is picked by spec "look" +(also --look and the MCP `look` argument), then env AINOTATE_LOOK, then `look` in config.toml, +then "default". A look that is not installed, fails to load, fails to start or fails while +drawing never breaks a render: the default look is used and the render gets a warning.""" +from __future__ import annotations + +import copy +import math +from types import SimpleNamespace + +from PIL import ImageColor, ImageDraw, ImageStat + +from .style import SHADOW, color_of, font_for, load_font, pill_size, rgba + + +def _r(): + """The render module: its drawing helpers are looked up at call time (tests patch them there).""" + from . import render + return render + + +class Look: + """render() draws every mark through these methods. `ctx` holds the unmarked `page`, `marks`, + `rects` and `points` (image px, after the crop), the crop's `offset` (image px), `scale`, `u`.""" + shadow = SHADOW # drop shadow under all marks: opacity, y-offset, blur (x u) + dim = 0.55 # spotlight darkness when the spec sets no "dim" + loupe_line = None # magnifier line width (None: the built-in width) + + def __init__(self, ctx): + self.ctx, u = ctx, ctx.u + self.stroke = max(3, round(u / 4)) + self.spot_radius = int(u) + self.font = load_font(int(1.35 * u)) + self.badge_font = load_font(int(1.4 * u)) + + def color(self, i): + return color_of(self.ctx.marks[i].get("color")) + + def label_size(self, text): + return pill_size(self.font, text, self.ctx.u) + + def label(self, layer, img, i, box, text): # labels and text notes; img: the page with tints + _r().draw_pill(layer, box, text, self.font, self.color(i), self.ctx.u) + + def outline(self, layer, i, g): + _r().draw_outline(layer, g, self.color(i), self.stroke, int(0.6 * self.ctx.u)) + + def badge(self, layer, img, i, b, n): + _r().draw_badge(layer, b, self.color(i), n, self.badge_font, max(2, round(0.12 * self.ctx.u))) + + def arrow(self, layer, gray, i, path, style, badge): # badge: the mark's own step badge or None + if style == "line": + _r().draw_arrow(ImageDraw.Draw(layer), path[0], path[-1], self.color(i), self.stroke, self.ctx.u) + else: + _r().draw_arrow_skitch(layer, gray, path[0], path[-1], self.color(i), self.ctx.u, path) + + def highlight(self, over, img, i, r): # lighter tint on dark regions (70 reads as a solid block there) + dark = ImageStat.Stat(img.convert("L").crop([int(v) for v in r])).mean[0] < 90 + ImageDraw.Draw(over).rectangle(r, fill=rgba(self.color(i), 45 if dark else 70)) + + def redact_fill(self, i): # a look may change the color of a redaction, never its area + return "#1F1F1F" + + def keys_font(self, i, keys): + return font_for(keys, int(1.15 * self.ctx.u)) + + +def installed_looks() -> dict: + """Looks from other packages: name -> entry point.""" + from importlib.metadata import entry_points + return {e.name: e for e in entry_points(group="ainotate.looks")} + + +class LookFailed(Exception): + """An installed look raised while drawing; render() then redraws with the default look.""" + + +def _check(ok, what, v): + if not ok: + raise ValueError(f"{what} is {v!r}") + return v + + +def _num(v, lo=0.0, hi=math.inf): + return isinstance(v, (int, float)) and not isinstance(v, bool) and math.isfinite(v) and lo <= v <= hi + + +def _color(v): + """The color as Pillow draws it ("info" -> "#007AFF"); ValueError for anything it cannot draw.""" + c = color_of(v) + ImageColor.getrgb(c) + return c + + +# what render() reads from a look, checked before use: a bad value is a failure of the look +CHECKS = { + "stroke": lambda v: _check(isinstance(v, int) and _num(v, 1), "stroke", v), + "spot_radius": lambda v: _check(isinstance(v, int) and _num(v), "spot_radius", v), + "dim": lambda v: _check(_num(v, 0, 1), "dim", v), + "loupe_line": lambda v: _check(v is None or (isinstance(v, int) and _num(v, 1)), "loupe_line", v), + "shadow": lambda v: _check(isinstance(v, tuple) and len(v) == 3 and all(_num(x) for x in v), "shadow", v), + "color": _color, + "redact_fill": _color, + "label_size": lambda v: _check(isinstance(v, tuple) and len(v) == 2 and all(_num(x, 1, 1e6) for x in v), + "label_size()", v), + "keys_font": lambda v: _check(hasattr(v, "getbbox") and hasattr(v, "getmetrics"), "keys_font()", v), +} + + +READS_PAGE = {"label", "badge", "highlight"} # hooks given the page image as their second argument + + +class _Guarded: + """An installed look whose errors and bad values become LookFailed (attribute reads too), so + they are never mistaken for core errors and render() can redraw with the default look.""" + + def __init__(self, look, name): + self._look, self._name = look, name + + def _fail(self, e): + return LookFailed(fallback(self._name, "failed while drawing", e)) + + def __getattr__(self, attr): + check = CHECKS.get(attr, lambda x: x) + try: + v = getattr(self._look, attr) + if not callable(v): + return check(v) + except Exception as e: # noqa: BLE001 - any failure inside a plug-in look + raise self._fail(e) from e + + def call(*a, **k): + if attr in READS_PAGE: # (layer, img, ...): drawing on its own copy of the page can never + a = (a[0], a[1].copy(), *a[2:]) # paint over a redaction + try: + return check(v(*a, **k)) + except Exception as e: # noqa: BLE001 + raise self._fail(e) from e + return call + + +def _private(ctx): + """A copy of ctx for an installed look: nothing it changes reaches the render or the spec.""" + return SimpleNamespace(page=ctx.page.copy(), marks=copy.deepcopy(ctx.marks), + rects=copy.deepcopy(ctx.rects), points=copy.deepcopy(ctx.points), + offset=ctx.offset, scale=ctx.scale, u=ctx.u) + + +def fallback(name, what, e=None) -> str: + why = f" ({type(e).__name__}: {e})" if e is not None else "" + return f"look {name!r} {what}{why}, so the default look was used" + + +def name_of(spec): + """The look a spec asks for: its "look", else AINOTATE_LOOK, else `look` in config.toml.""" + if "look" in spec: + return spec["look"] + from .output import default_look + return default_look() + + +def installed_name(spec): + """The name of the installed look this spec would use, else None (never raises).""" + try: + name = name_of(spec) + return name if name != "default" and name in installed_looks() else None + except Exception: # noqa: BLE001 + return None + + +def pick(spec, ctx, warnings): + """The look for this render; a look that cannot be used adds a warning and gives the default.""" + name = name_of(spec) + if name == "default": + return Look(ctx) + found = installed_looks() + if name not in found: + warnings.append(fallback(name, "is not installed") + "; installed looks: " + + ", ".join(["default", *sorted(found)])) + return Look(ctx) + try: + cls = found[name].load() + except Exception as e: # noqa: BLE001 + warnings.append(fallback(name, "could not be loaded", e)) + return Look(ctx) + try: + return _Guarded(cls(_private(ctx)), name) + except Exception as e: # noqa: BLE001 + warnings.append(fallback(name, "failed to start", e)) + return Look(ctx) diff --git a/src/ainotate/marks.py b/src/ainotate/marks.py index eb1098a..5f2c257 100644 --- a/src/ainotate/marks.py +++ b/src/ainotate/marks.py @@ -92,12 +92,14 @@ def _shape(d, box, shape, radius, **kw): d.rounded_rectangle(box, radius=radius, **kw) -def draw_loupe(layer, page, sample, box, shape, color, u, path=None): +def draw_loupe(layer, page, sample, box, shape, color, u, path=None, line=None): """Source outline, connector and the loupe itself (zoomed page pixels, white rim, color ring). `page` is the image before any mark is drawn (redactions and blur already applied). `path`: a - curved connector from the loupe to the sample (routed around page content), else straight.""" + curved connector from the loupe to the sample (routed around page content), else straight. + `line`: width of the connector, source outline and color ring (default from u).""" lw, lh = int(box[2] - box[0]), int(box[3] - box[1]) - line = max(2, round(0.13 * u)) + ring = line + line = line or max(2, round(0.13 * u)) rad_src = min(0.6 * u, (sample[3] - sample[1]) / 2) rad_lp = min(1.3 * u, min(lw, lh) / 4) # connector between the facing edges (drawn first: the loupe and outline cover its ends) @@ -136,7 +138,7 @@ def draw_loupe(layer, page, sample, box, shape, color, u, path=None): mask = _aa_mask((lw, lh), lambda d, S: _shape(d, [0, 0, lw * S - 1, lh * S - 1], shape, rad_lp * S, fill=255)) content.putalpha(mask) layer.alpha_composite(content, dest=(int(box[0]), int(box[1]))) - rim, ring = max(3, round(0.3 * u)), max(2, round(0.1 * u)) + rim, ring = max(3, round(0.3 * u)), ring or max(2, round(0.1 * u)) inner = grow(box, -rim / 2) aa_draw(layer, box, "white", lambda d, T, S: _shape(d, T(inner), shape, (rad_lp - rim / 2) * S, outline="white", width=rim * S)) diff --git a/src/ainotate/mcp_docs.py b/src/ainotate/mcp_docs.py index 6512f11..0e58b21 100644 --- a/src/ainotate/mcp_docs.py +++ b/src/ainotate/mcp_docs.py @@ -250,7 +250,8 @@ def _fix_for(e: BaseException) -> str: _FALLBACK_SPEC = """Spec: {"input": image path, "scale": spec units -> image px, "crop": "auto" | -[x1,y1,x2,y2] | {x,y,w,h}, "crop_pad": 260, "dim": 0.55, "name": "file name", "marks": [...]} +[x1,y1,x2,y2] | {x,y,w,h}, "crop_pad": 260, "dim": 0.55, "name": "file name", "look": "default" or an +installed look such as "neat" (mark style), "marks": [...]} Units: every coordinate (rect, crop, label_at, at, pad) in ONE unit. DOM rects from capture_web: dicts {x,y,w,h} with the returned scale. locate / grid / zoom pixels: lists [x1,y1,x2,y2], scale 1. Colors: look (orange, default), bad (red: only for wrong/bug), good (green), info (blue). diff --git a/src/ainotate/mcp_server.py b/src/ainotate/mcp_server.py index f221885..affa059 100644 --- a/src/ainotate/mcp_server.py +++ b/src/ainotate/mcp_server.py @@ -12,10 +12,11 @@ import json import sys from pathlib import Path -from typing import Callable, Optional, Union +from typing import Annotated, Callable, Optional, Union import anyio.to_thread from mcp import types +from pydantic import Field from PIL import Image from . import __version__ @@ -94,6 +95,17 @@ def _spec(spec) -> dict: return spec +LOOK = Annotated[Optional[str], Field(description=( + 'Mark style: "default" (bold, built in) or the name of an installed look package (e.g. "neat"). ' + "Omit it to use the spec's look, else the user's default (AINOTATE_LOOK, config.toml). A look " + "that is not installed or fails falls back to default with a warning listing the installed looks."))] + + +def _with_look(spec, look): + spec = _spec(spec) + return dict(spec, look=look) if look and isinstance(spec, dict) else spec + + # ---------------------------------------------------------------- annotate def _annotate(spec, draft: bool, debug: bool, tool: str) -> types.CallToolResult: @@ -120,22 +132,23 @@ def _annotate(spec, draft: bool, debug: bool, tool: str) -> types.CallToolResult @mcp.tool() -async def annotate(spec: Union[dict, str], draft: bool = True) -> types.CallToolResult: +async def annotate(spec: Union[dict, str], draft: bool = True, look: LOOK = None) -> types.CallToolResult: """Render an annotation spec onto its `input` image. Draft by default (temp dir): check the returned preview, then call again with draft=false to save to the output folder. spec: {"input": path, "scale": 1|2.., "crop": "auto"|[x1,y1,x2,y2], "name": "...", "marks": [{"type": "step"|"box"|"arrow"|"highlight"|"spotlight"|"redact"|"text"|..., "rect": {x,y,w,h} or [x1,y1,x2,y2], "label": "...", "n": 1, "color": "look|bad|good|info"}]}. Every coordinate in one unit; scale converts it to image px. Full reference: ainotate://spec. + look: the mark style, "default" or an installed look such as "neat" (sets spec "look"). Returns paths, size, warnings, legend (step numbers + labels), redactions and a preview.""" - return await _run("annotate", lambda: _annotate(spec, draft, False, "annotate")) + return await _run("annotate", lambda: _annotate(_with_look(spec, look), draft, False, "annotate")) @mcp.tool() -async def preview(spec: Union[dict, str]) -> types.CallToolResult: +async def preview(spec: Union[dict, str], look: LOOK = None) -> types.CallToolResult: """Draft render with the debug overlay (cyan target rects, magenta label boxes); saves - nothing permanent. Use it to check placement before annotate.""" - return await _run("preview", lambda: _annotate(spec, True, True, "preview")) + nothing permanent. Use it to check placement before annotate. look: as in annotate.""" + return await _run("preview", lambda: _annotate(_with_look(spec, look), True, True, "preview")) # ---------------------------------------------------------------- web @@ -180,12 +193,12 @@ def work(): @mcp.tool() -async def shoot(spec: Union[dict, str], draft: bool = True) -> types.CallToolResult: +async def shoot(spec: Union[dict, str], draft: bool = True, look: LOOK = None) -> types.CallToolResult: """One call from URL to annotated image: spec = render spec without input/scale, plus url (or cdp_url + page_url_contains), preset, actions, and marks with "target": {...} (same forms as capture_web) instead of "rect". Draft by default: check the preview, then call again with draft=false to save. Returns paths, warnings, legend, redactions (secrets - hidden automatically), the capture's resolved rects and a preview.""" + hidden automatically), the capture's resolved rects and a preview. look: as in annotate.""" def work(): try: from .shoot import shoot as _shoot @@ -196,7 +209,7 @@ def work(): "fix": "Use capture_web(targets=...) and then annotate(spec) with the returned " "rects and scale."}}, error=True) raise - res = _shoot(_spec(spec), draft=draft) + res = _shoot(_with_look(spec, look), draft=draft) paths = [str(p) for p in res.paths] jpeg, (w, h) = preview_jpeg(paths[0]) cap = getattr(res, "capture", None) diff --git a/src/ainotate/output.py b/src/ainotate/output.py index 5270d20..f08b0b6 100644 --- a/src/ainotate/output.py +++ b/src/ainotate/output.py @@ -60,6 +60,14 @@ def pick(env, key, default): return Config(Path(out).expanduser(), Path(bak).expanduser() if bak else None, prefix) +def default_look() -> str: + """The look for a spec without "look": env AINOTATE_LOOK > config.toml `look` > "default".""" + look = os.environ.get("AINOTATE_LOOK") + if not look: + look = _read_config_file(CONFIG_PATH).get("look") + return look if isinstance(look, str) and look else "default" + + def safe_component(value: str, what: str = "name") -> str: """`value` if it is one plain file-name component: no path separators, no '..', no control characters, not '.'; else OutputError (a bad prefix must not write outside the folder).""" diff --git a/src/ainotate/render.py b/src/ainotate/render.py index 00fda57..13688d1 100644 --- a/src/ainotate/render.py +++ b/src/ainotate/render.py @@ -27,8 +27,12 @@ OCR on the input). spec "privacy": "auto" adds solid redactions for secrets OCR reads (default "off" here); spec "frame" puts the finished image on a backdrop (beautify.apply_frame, last step). +spec "look": the mark style, "default" (built in) or the name of an installed look package (entry +point group "ainotate.looks", e.g. "neat"); without it AINOTATE_LOOK, then `look` in config.toml. + Warnings: more than 6 marks, labels over 4 words, labels that overlap each other or cover another -mark's target, arrows that cross, a magnifier that had to shrink its zoom. +mark's target, arrows that cross, a magnifier that had to shrink its zoom, a look that is not +installed or failed (the default look is used then; the warning lists the installed looks). Layout (crop, labels, badges, loupes) lives in placement.py. """ @@ -36,6 +40,7 @@ import math from dataclasses import dataclass, field +from types import SimpleNamespace from PIL import Image, ImageDraw, ImageFilter, ImageStat @@ -44,6 +49,7 @@ overlap_warnings, route_arrow) from .label_plan import Job, plan_labels from .placement import RenderError, ContentMap, auto_crop, place_badge, place_label, place_loupe # noqa: F401 +from .looks import Look, LookFailed, fallback, installed_name, pick # noqa: F401 (Look: base class) from .spec import DECORATIVE, LABELED, open_image, spec_warnings, validate from .style import LABEL_FILL, SHADOW, color_of, font_for, load_font, mark_unit, pill_size, rgba, text_color_on @@ -133,7 +139,24 @@ def _outside_crop(i, m, r, crop, scale): def render(spec: dict, debug: bool = False) -> RenderResult: """Validate and draw a spec. SpecError for an invalid spec or unreadable input, - RenderError for a valid spec that cannot be drawn.""" + RenderError for a valid spec that cannot be drawn. When an installed look is in use and + anything fails (the look, or AInotate on a value the look gave it), the render is redone from + the spec with the default look, with a warning; an error the default look also hits is raised.""" + try: + return _render(spec, debug) + except LookFailed as e: + warning = str(e) + except Exception as e: # noqa: BLE001 - only with an installed look, see below + name = installed_name(spec) + if name is None: + raise + warning = fallback(name, "failed while drawing", e) + res = _render(dict(spec, look="default"), debug) + res.warnings.insert(0, warning) + return res + + +def _render(spec: dict, debug: bool) -> RenderResult: from .shoot import prepare_image_spec # text targets (OCR) and auto privacy -> rects n_user = len(validate(spec)) spec, redactions = prepare_image_spec(spec) @@ -197,9 +220,9 @@ def render(spec: dict, debug: bool = False) -> RenderResult: W, H = img.size u = mark_unit(W, H, scale) style = spec.get("arrow_style", "skitch") - stroke = max(3, round(u / 4)) - font = load_font(int(1.35 * u)) - badge_font = load_font(int(1.4 * u)) + lk = pick(spec, SimpleNamespace(page=page, marks=marks, rects=rects, points=points, + offset=(ox, oy), scale=scale, u=u), warnings) + stroke = lk.stroke base_pad = 0.7 * u def pad_for(m, r): @@ -216,26 +239,24 @@ def pad_for(m, r): key, default = ("radius", 0.5 * u) if m["type"] == "blur" else ("block", 0.9 * u) img = deemphasize(img, r, m["type"], m[key] * scale if key in m else default, u) over = Image.new("RGBA", img.size, (0, 0, 0, 0)) - od = ImageDraw.Draw(over) spots = [g for m, g in zip(marks, outlines) if m["type"] == "spotlight"] if spots: - dim = Image.new("RGBA", img.size, (0, 0, 0, int(255 * spec.get("dim", 0.55)))) + dim = Image.new("RGBA", img.size, (0, 0, 0, int(255 * spec.get("dim", lk.dim)))) mask = Image.new("L", img.size, 255) md = ImageDraw.Draw(mask) for g in spots: - md.rounded_rectangle(g, radius=int(u), fill=0) + md.rounded_rectangle(g, radius=lk.spot_radius, fill=0) over.paste(dim, (0, 0), mask) - for m, r in zip(marks, rects): - if m["type"] == "highlight": # lighter tint on dark regions (70 reads as a solid block there) - dark = ImageStat.Stat(img.convert("L").crop([int(v) for v in r])).mean[0] < 90 - od.rectangle(r, fill=rgba(color_of(m.get("color")), 45 if dark else 70)) + for i, (m, r) in enumerate(zip(marks, rects)): + if m["type"] == "highlight": + lk.highlight(over, img, i, r) img = Image.alpha_composite(img, over) d = ImageDraw.Draw(img) - for m, r in zip(marks, rects): + for i, (m, r) in enumerate(zip(marks, rects)): if m["type"] == "redact": # solid, never blur: blur can be reversed out = m["pad"] * scale if "pad" in m else max(2, 2 * scale) # default outset: safer - d.rectangle(grow(r, out), fill=rgba(color_of(m.get("fill", "#1F1F1F")), 255)) + d.rectangle(grow(r, out), fill=rgba(color_of(m.get("fill", lk.redact_fill(i))), 255)) # 2) shapes; collect obstacles so labels avoid them # placement reads page content (edges = text, lines, icons) before any mark covers it @@ -247,10 +268,10 @@ def pad_for(m, r): d = ImageDraw.Draw(layer) obstacles, placed, dbg, badges = [], [], [], [] targets = [r for r, a in zip(rects, auto) if r and not a] # highlight/spotlight/redact/arrow too - for m, r, g in zip(marks, rects, outlines): + for i, (m, r, g) in enumerate(zip(marks, rects, outlines)): if m["type"] in ("box", "step"): if m.get("box", True): - draw_outline(layer, g, color_of(m.get("color")), stroke, int(0.6 * u)) + lk.outline(layer, i, g) obstacles.append(g) # forced label positions are known up front: reserve them so no badge or auto label lands there @@ -258,7 +279,7 @@ def pad_for(m, r): for i, m in enumerate(marks): if m["type"] in LABELED and m.get("label_at"): lx, ly = [v * scale for v in m["label_at"]] - pw, ph = pill_size(font, m["label"], u) + pw, ph = lk.label_size(m["label"]) forced[i] = clamp_box([lx - ox, ly - oy, lx - ox + pw, ly - oy + ph], W, H) if forced[i]: placed.append(forced[i]) @@ -266,11 +287,11 @@ def pad_for(m, r): # click / keys-at / magnify have fixed or early positions: lay them out before labels and badges own, hard, labels, arrows = {}, [], [], [] # own: box a mark's label and overlap checks refer to leaders = [] # loupe connectors (labels avoid them like arrows) - key_fonts = {i: font_for("".join(m["keys"]), int(1.15 * u)) for i, m in enumerate(marks) if m["type"] == "keys"} + key_fonts = {i: lk.keys_font(i, "".join(m["keys"])) for i, m in enumerate(marks) if m["type"] == "keys"} for i, (m, r) in enumerate(zip(marks, rects)): if m["type"] == "click": x, y = points.get(i) or ((r[0] + r[2]) / 2, (r[1] + r[3]) / 2) - draw_click(layer, x, y, color_of(m.get("color")), u) + draw_click(layer, x, y, lk.color(i), u) own[i] = click_geometry(x, y, u)[2] obstacles.append(own[i]) hard.append(own[i]) @@ -289,7 +310,7 @@ def pad_for(m, r): avoid = [o for o in obstacles + targets if o != r] box, sample, path = place_loupe(i, m, r, points.get(i), W, H, u, cmap, avoid, [o for o in targets + hard if o != r], placed, warnings, leaders) - draw_loupe(layer, img, sample, box, m.get("shape", "circle"), color_of(m.get("color")), u, path) + draw_loupe(layer, img, sample, box, m.get("shape", "circle"), lk.color(i), u, path, lk.loupe_line) if path: leaders.append(path) own[i] = box @@ -309,7 +330,7 @@ def pad_for(m, r): if not clean: warnings.append(f"mark #{i}: badge overlaps the target (target touches the frame edge); " "widen the crop") - badges.append((b, color_of(m.get("color")), str(m["n"]))) # drawn last, above arrows and labels + badges.append((b, i, str(m["n"]))) # drawn last, above arrows and labels badge_of[i] = b placed.append(b) @@ -319,7 +340,7 @@ def pad_for(m, r): if m["type"] in LABELED and m.get("label") and not m.get("label_at"): target = own.get(i) or (g if m["type"] != "arrow" else r) near = 5 if m["type"] == "arrow" else (3.5 if m.get("arrow", True) else 2) - jobs.append(Job(i, target, pill_size(font, m["label"], u), + jobs.append(Job(i, target, lk.label_size(m["label"]), [o for o in obstacles + targets if o != target and o != r], [o for o in targets + hard if o != target and o != r], near, dict(arrow=m.get("arrow", True), curve=style in ("skitch", "curved"), style=style, @@ -334,7 +355,7 @@ def pad_for(m, r): "shorten the label or set label_at") for i, (m, r, g) in enumerate(zip(marks, rects, outlines)): - t, c = m["type"], color_of(m.get("color")) + t = m["type"] if t in LABELED and m.get("label"): target = own.get(i) or (g if t != "arrow" else r) lb = forced.get(i) @@ -351,11 +372,8 @@ def pad_for(m, r): path = route_arrow(edges, lb, target, u, avoid, style, stroke, [a[3] for a in arrows] + leaders) if path_length(path) > 2.8 * u: # room for a body behind the head arrows.append((i, path[0], path[-1], path)) - if style == "line": - draw_arrow(d, path[0], path[-1], c, stroke, u) - else: - draw_arrow_skitch(layer, gray, path[0], path[-1], c, u, path) - draw_pill(layer, lb, m["label"], font, c, u) + lk.arrow(layer, gray, i, path, style, badge_of.get(i)) + lk.label(layer, img, i, lb, m["label"]) dbg.append(lb) labels.append((i, lb)) if t == "keys" and i not in points: # beside its rect, like a label without an arrow @@ -373,21 +391,21 @@ def pad_for(m, r): labels.append((i, kb)) if t == "text": x, y = [v * scale for v in m["at"]] - pw, ph = pill_size(font, m["text"], u) + pw, ph = lk.label_size(m["text"]) tb = clamp_box([x - ox, y - oy, x - ox + pw, y - oy + ph], W, H) if tb is None: raise RenderError(f"mark #{i} (text): note {m['text']!r} is larger than the frame; widen the " "crop or shorten it") - draw_pill(layer, tb, m["text"], font, c, u) + lk.label(layer, img, i, tb, m["text"]) placed.append(tb) dbg.append(tb) labels.append((i, tb)) - for b, c, n in badges: - draw_badge(layer, b, c, n, badge_font, max(2, round(0.12 * u))) + for b, i, n in badges: + lk.badge(layer, img, i, b, n) # Skitch-style soft drop shadow under all marks, then the marks themselves - opacity, dy, blur = SHADOW + opacity, dy, blur = lk.shadow alpha = layer.getchannel("A").filter(ImageFilter.GaussianBlur(blur * u)) shadow = Image.new("RGBA", img.size, (0, 0, 0, 0)) shadow.paste(Image.new("RGBA", img.size, (0, 0, 0, 255)), (0, round(dy * u)), diff --git a/src/ainotate/spec.py b/src/ainotate/spec.py index 83eae4e..28addfd 100644 --- a/src/ainotate/spec.py +++ b/src/ainotate/spec.py @@ -100,6 +100,8 @@ def validate(spec): errs.append("spec: 'crop_pad' must be a number >= 0") if spec.get("arrow_style", "skitch") not in ("skitch", "straight", "curved", "line"): errs.append("spec: 'arrow_style' must be skitch, straight, curved or line") + if "look" in spec and not (isinstance(spec["look"], str) and spec["look"]): + errs.append("spec: 'look' must be a look name (\"default\" or an installed look)") if "name" in spec and not isinstance(spec["name"], str): errs.append("spec: 'name' must be a string") if "marks" not in spec: diff --git a/tests/conftest.py b/tests/conftest.py index 3315cc6..6802f69 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -43,7 +43,8 @@ def isolated_env(tmp_path, monkeypatch): home.mkdir() monkeypatch.setenv("HOME", str(home)) monkeypatch.setenv("USERPROFILE", str(home)) - for v in ("AINOTATE_OUTPUT_DIR", "AINOTATE_BACKUP_DIR", "AINOTATE_PREFIX", "AINOTATE_FONT", "AINOTATE_NEO_DIR"): + for v in ("AINOTATE_OUTPUT_DIR", "AINOTATE_BACKUP_DIR", "AINOTATE_PREFIX", "AINOTATE_FONT", "AINOTATE_NEO_DIR", + "AINOTATE_LOOK"): monkeypatch.delenv(v, raising=False) return home diff --git a/tests/test_looks.py b/tests/test_looks.py new file mode 100644 index 0000000..d4ee618 --- /dev/null +++ b/tests/test_looks.py @@ -0,0 +1,294 @@ +import hashlib +import importlib.metadata +import json + +import PIL +import pytest +from PIL import Image, features + +from ainotate import cli +from ainotate.render import Look, render +from ainotate.spec import SpecError + +EVERY = [ + {"type": "step", "n": 1, "rect": [1300, 100, 1520, 150], "label": "Checkout"}, + {"type": "box", "rect": [1300, 170, 1520, 220], "label": "Coupon", "color": "good"}, + {"type": "arrow", "rect": [280, 450, 900, 480], "label": "Row", "color": "bad"}, + {"type": "highlight", "rect": [280, 340, 900, 370]}, + {"type": "spotlight", "rect": [280, 560, 900, 600]}, + {"type": "redact", "rect": {"x": 420, "y": 230, "w": 170, "h": 14}}, + {"type": "text", "text": "A note", "at": [700, 900], "color": "info"}, + {"type": "magnify", "rect": [1300, 100, 1400, 150], "at": [1100, 800]}, + {"type": "click", "at": [1410, 125], "label": "Click", "color": "info"}, + {"type": "keys", "keys": ["Ctrl", "S"], "at": [300, 940]}, + {"type": "blur", "rect": [290, 680, 700, 700]}, + {"type": "pixelate", "rect": [290, 735, 700, 755]}, +] +# sha256 of the default render of every mark type, made with the code before looks were pluggable +DEFAULT_SHA = {False: "2459af01e003be489fe9675a07cea01f3109a6ed153d6b02d767dd0d11456417", + True: "fc93f1192da4bbcc2fb0ad05d98fa91b006776c58c0968f158ffa25769628368"} + + +@pytest.mark.parametrize("dark", [False, True]) +def test_default_look_is_unchanged(ui, ui_dark, dark): + s = {"input": ui_dark if dark else ui, "marks": EVERY} + img = render(s).image + assert render(dict(s, look="default")).image.tobytes() == img.tobytes() + if (PIL.__version__, features.version("freetype2")) != ("12.3.0", "2.14.3"): + pytest.skip("pinned pixels were made with Pillow 12.3.0 and FreeType 2.14.3") + assert hashlib.sha256(img.tobytes()).hexdigest() == DEFAULT_SHA[dark] + + +class Gray(Look): + """A test look: gray outlines and redactions, no shadow; records what it drew.""" + shadow = (0, 0, 0) + drawn = [] + + def outline(self, layer, i, g): + Gray.drawn.append(("outline", i)) + super().outline(layer, i, g) + + def color(self, i): + return "#777777" + + def redact_fill(self, i): + return "#555555" + + +class EP: + name, group = "gray", "ainotate.looks" + + def load(self): + return Gray + + +@pytest.fixture +def gray_installed(monkeypatch): + real = importlib.metadata.entry_points + monkeypatch.setattr(importlib.metadata, "entry_points", + lambda **kw: [EP()] if kw.get("group") == "ainotate.looks" else real(**kw)) + Gray.drawn = [] + + +def test_an_installed_look_draws_the_marks(ui, gray_installed): + marks = [EVERY[1], EVERY[5]] + plain = render({"input": ui, "marks": marks}).image + img = render({"input": ui, "look": "gray", "marks": marks}).image + assert img.tobytes() != plain.tobytes() and Gray.drawn == [("outline", 0)] + assert img.getpixel((500, 237)) == (0x55, 0x55, 0x55) # the redaction, in the look's color + + +def test_look_from_flag_env_and_config(ui, gray_installed, isolated_env, monkeypatch, tmp_path, capsys): + spec = tmp_path / "s.json" + spec.write_text(json.dumps({"input": ui, "marks": [EVERY[1]]})) + assert cli.main(["annotate", str(spec), "--draft", "--look", "gray"]) == 0 and Gray.drawn + cfg = isolated_env / ".config" / "ainotate" / "config.toml" + cfg.parent.mkdir(parents=True) + cfg.write_text('look = "gray"\n') + for env, used in ((None, True), ("default", False), ("gray", True)): + Gray.drawn = [] + if env: + monkeypatch.setenv("AINOTATE_LOOK", env) + render({"input": ui, "marks": [EVERY[1]]}) + assert bool(Gray.drawn) == used, env + Gray.drawn = [] + render({"input": ui, "look": "default", "marks": [EVERY[1]]}) # the spec beats env and config + assert not Gray.drawn + + +@pytest.mark.parametrize("where", ["spec", "env", "config"]) +def test_a_look_that_is_not_installed_falls_back_with_a_warning(ui, gray_installed, isolated_env, monkeypatch, where): + s = {"input": ui, "marks": [EVERY[1]]} + if where == "spec": + s["look"] = "neat" + elif where == "env": + monkeypatch.setenv("AINOTATE_LOOK", "neat") + else: + cfg = isolated_env / ".config" / "ainotate" / "config.toml" + cfg.parent.mkdir(parents=True) + cfg.write_text('look = "neat"\n') + res = render(s) + assert res.image.tobytes() == render({"input": ui, "marks": [EVERY[1]]}).image.tobytes() + assert res.warnings == ["look 'neat' is not installed, so the default look was used; " + "installed looks: default, gray"] + + +class Broken(Gray): + def __init__(self, ctx): + raise RuntimeError("no font") + + +class BadDraw(Gray): + def label(self, layer, img, i, box, text): + raise ValueError("bad label") + + +@pytest.mark.parametrize("case, warning", [ + ("import", "look 'gray' could not be loaded (ModuleNotFoundError: No module named 'gray_look'), " + "so the default look was used"), + ("init", "look 'gray' failed to start (RuntimeError: no font), so the default look was used"), + ("draw", "look 'gray' failed while drawing (ValueError: bad label), so the default look was used"), +]) +def test_a_broken_look_falls_back_with_a_warning(ui, gray_installed, monkeypatch, tmp_path, capsys, case, warning): + def load(): + if case == "import": + import gray_look # noqa: F401 + return Broken if case == "init" else BadDraw + monkeypatch.setattr(EP, "load", lambda self: load()) + s = {"input": ui, "look": "gray", "marks": [EVERY[0], EVERY[5]]} + res = render(s) + assert res.warnings[0] == warning + assert res.image.tobytes() == render(dict(s, look="default")).image.tobytes() + spec = tmp_path / "s.json" # the CLI prints it like any render warning + spec.write_text(json.dumps(s)) + assert cli.main(["annotate", str(spec), "--draft"]) == 0 + assert warning in capsys.readouterr().err + + +def test_mcp_tools_take_a_look_and_return_its_warning(ui, gray_installed, out_dirs): + pytest.importorskip("mcp") + import anyio + from mcp import Client + + from ainotate.mcp_server import mcp + + async def go(): + async with Client(mcp) as c: + tools = {t.name: t for t in (await c.list_tools()).tools} + for name in ("annotate", "preview", "shoot"): + assert "installed look" in tools[name].input_schema["properties"]["look"]["description"] + res = await c.call_tool("annotate", {"spec": {"input": ui, "marks": [EVERY[1]]}, "look": "nope"}) + return json.loads(res.content[0].text) + d = anyio.run(go) + assert d["ok"] and d["warnings"] == ["look 'nope' is not installed, so the default look was used; " + "installed looks: default, gray"] + + +def test_look_must_be_a_name(ui): + with pytest.raises(SpecError) as e: + render({"input": ui, "look": 3, "marks": []}) + assert "'look' must be a look name" in str(e.value) + + +# ---------------------------------------------------------------- a look can never change redactions + +REDACT = {"type": "redact", "rect": [100, 100, 300, 130]} + + +class Mutates(Gray): + """Changes the geometry and spec data it is given, then fails (in __init__ or in outline).""" + when = "init" + + def __init__(self, ctx): + super().__init__(ctx) + ctx.rects[0][0] += 10 # the redaction's left edge + ctx.marks[1].pop("label", None) + if Mutates.when == "init": + raise RuntimeError("bad start") + + def outline(self, layer, i, g): + raise RuntimeError("bad outline") + + +class Paints(Gray): + """Works, but paints white over the whole page image it is handed.""" + + def label(self, layer, img, i, box, text): + img.paste((255, 255, 255, 255), (0, 0, img.width, img.height)) + super().label(layer, img, i, box, text) + + +@pytest.mark.parametrize("cls, when", [(Mutates, "init"), (Mutates, "draw"), (Paints, None)]) +def test_a_look_never_changes_redaction_coverage_or_the_spec(ui, gray_installed, monkeypatch, cls, when): + import copy + Mutates.when = when + monkeypatch.setattr(EP, "load", lambda self: cls) + s = {"input": ui, "look": "gray", "marks": [REDACT, EVERY[1]]} + before = copy.deepcopy(s) + res = render(s) + assert s == before + assert res.image.getpixel((100, 110)) == res.image.getpixel((299, 129)) != (255, 255, 255) + plain = render(dict(s, look="default")).image + if cls is Mutates: # the failed render is redone from the spec, exactly as the default look + assert res.image.tobytes() == plain.tobytes() + assert res.warnings[0].startswith("look 'gray' failed") + else: # a working look: drawn by it, on its own copy of the page + assert not res.warnings and res.image.tobytes() != plain.tobytes() + + +class BadValues(Gray): + bad = "shadow" + + @property + def shadow(self): + if BadValues.bad == "shadow": + raise RuntimeError("no shadow") + return (0, 0, 0) + + def redact_fill(self, i): + return "not-a-color" if BadValues.bad == "redact_fill" else "#555555" + + def label_size(self, text): + return (0, -1) if BadValues.bad == "label_size" else super().label_size(text) + + +@pytest.mark.parametrize("bad, why", [ + ("shadow", "RuntimeError: no shadow"), + ("redact_fill", "ValueError: unknown color specifier: 'not-a-color'"), + ("label_size", "ValueError: label_size() is (0, -1)"), +]) +def test_bad_attributes_and_hook_results_fall_back(ui, gray_installed, monkeypatch, tmp_path, capsys, bad, why): + BadValues.bad = bad + monkeypatch.setattr(EP, "load", lambda self: BadValues) + s = {"input": ui, "look": "gray", "marks": [REDACT, EVERY[1]]} + res = render(s) + assert res.warnings[0] == f"look 'gray' failed while drawing ({why}), so the default look was used" + assert res.image.tobytes() == render(dict(s, look="default")).image.tobytes() + spec = tmp_path / "s.json" + spec.write_text(json.dumps(s)) + assert cli.main(["annotate", str(spec), "--draft"]) == 0 + assert why in capsys.readouterr().err + + +# ---------------------------------------------------------------- values the core cannot use + +class Odd(Gray): + kind = "name" + + def color(self, i): + return "info" if Odd.kind == "name" else super().color(i) + + @property + def shadow(self): + return (0.3, float("inf"), 0.15) if Odd.kind == "inf" else (0, 0, 0) + + def label_size(self, text): # passes the checks, but no frame has room for it + return (999_999, 999_999) if Odd.kind == "huge" else super().label_size(text) + + +def test_a_color_name_from_a_look_is_drawn_as_its_color(ui, gray_installed, monkeypatch): + Odd.kind = "name" + monkeypatch.setattr(EP, "load", lambda self: Odd) + click = {"type": "click", "at": [1410, 125]} # drawn by AInotate in the look's color(i) + res = render({"input": ui, "look": "gray", "marks": [click]}) + assert not res.warnings and res.image.getpixel((1410, 125)) != Image.open(ui).getpixel((1410, 125)) + + +@pytest.mark.parametrize("kind, why", [ + ("inf", "(ValueError: shadow is (0.3, inf, 0.15))"), + ("huge", "(RenderError: mark #0 (box): no room for label 'Coupon'."), +]) +def test_values_the_core_cannot_use_fall_back(ui, gray_installed, monkeypatch, kind, why): + Odd.kind = kind + monkeypatch.setattr(EP, "load", lambda self: Odd) + s = {"input": ui, "look": "gray", "marks": [EVERY[1]]} + res = render(s) + assert res.warnings[0].startswith("look 'gray' failed while drawing " + why) + assert res.warnings[0].endswith(", so the default look was used") + assert res.image.tobytes() == render(dict(s, look="default")).image.tobytes() + + +def test_an_error_the_default_look_also_hits_is_raised(ui, gray_installed): + from ainotate.render import RenderError + with pytest.raises(RenderError): # a label position outside the image is the spec's fault + render({"input": ui, "look": "gray", "marks": [dict(EVERY[1], label_at=[5000, 10])]})