plot: Support adding custom legend items

fixes #9

src/plot.typ

@@ -12,6 +12,7 @@
 #import "/src/plot/bar.typ": add-bar
 #import "/src/plot/errorbar.typ": add-errorbar
 #import "/src/plot/mark.typ"
+#import plot-legend: add-legend
 
 #let default-colors = (blue, red, green, yellow, black)
 
@@ -261,6 +262,8 @@
   // Create axes for data & annotations
   let axis-dict = (:)
   for d in data + annotations {
+    if "axes" not in d { continue }
+
     for (i, name) in d.axes.enumerate() {
       if not name in axis-dict {
         axis-dict.insert(name, axes.axis(
@@ -304,6 +307,8 @@
 
   // Prepare styles
   for i in range(data.len()) {
+    if "style" not in data.at(i) { continue }
+
     let style-base = plot-style
     if type(style-base) == function {
       style-base = (style-base)(i)
@@ -346,6 +351,8 @@
 
     // Prepare
     for i in range(data.len()) {
+      if "axes" not in data.at(i) { continue }
+
       let (x, y) = data.at(i).axes.map(name => axis-dict.at(name))
       let plot-ctx = make-ctx(x, y, size)
 
@@ -370,6 +377,8 @@
     // Fill
     if fill-below {
       for d in data {
+        if "axes" not in d { continue }
+
         let (x, y) = d.axes.map(name => axis-dict.at(name))
         let plot-ctx = make-ctx(x, y, size)
 
@@ -420,6 +429,8 @@
 
     // Stroke + Mark data
     for d in data {
+      if "axes" not in d { continue }
+
       let (x, y) = d.axes.map(name => axis-dict.at(name))
       let plot-ctx = make-ctx(x, y, size)
 
@@ -487,7 +498,7 @@
             mark: item.at("mark", default: none),
             mark-size: item.at("mark-size", default: none),
             mark-style: item.at("mark-style", default: none),
-            ..item.style)
+            ..item.at("style", default: (:)))
         }
       }, ..legend-style)
     }

src/plot/legend.typ

@@ -203,3 +203,30 @@
               "items.north-east", fill: style.fill, stroke: style.stroke, radius: style.radius)
   })
 })
+
+/// Function for manually adding a legend item from within
+/// a plot environment
+///
+/// - label (content): Legend label
+/// - preview (auto,function): Legend preview function of the format `() => elements`.
+///     The preview canvas bounds are between (0,0) and (1,1).
+///     If set to `auto`, a straight line is drawn.
+///
+/// ```example
+/// add-legend([Custom item], preview _ => {
+///   draw.rect((0,0), (1,1)) // Draw a rect as preview
+/// })
+/// ```
+#let add-legend(label, preview: auto) = {
+  assert(preview == auto or type(preview) == function,
+    message: "Expected auto or function, got " + repr(type(preview)))
+
+  return ((
+    type: "legend-item",
+    label: label,
+    style: (:),
+    axes: ("x", "y"),
+  ) + if preview != auto {
+    (plot-legend-preview: _ => { preview() })
+  },)
+}

tests/plot/legend/test.typ

@@ -1,15 +1,14 @@
 #set page(width: auto, height: auto)
-#import "/src/lib.typ": *
-#import "/src/cetz.typ": *
 #import "/tests/helper.typ": *
+#import cetz: draw
+#import cetz-plot: plot
 
 #let dom = (domain: (0, 2 * calc.pi))
 #let fn(x, offset: 0) = {calc.sin(x) + offset}
 
 #for pos in ("north", "south", "west", "east",
              "north-east", "north-west",
              "south-east", "south-west",) {
-  pos = "legend." + pos
   test-case({
     import draw: *
 
@@ -26,7 +25,6 @@
 #for pos in ("inner-north", "inner-south", "inner-west", "inner-east",
              "inner-north-east", "inner-north-west",
              "inner-south-east", "inner-south-west",) {
-  pos = "legend." + pos
   test-case({
     import draw: *
 
@@ -117,7 +115,7 @@
   import draw: *
 
   set-style(legend: (item: (preview: (width: .4), spacing: .7),
-    orientation: ltr, default-position: "legend.north"))
+    orientation: ltr, default-position: "north"))
 
   plot.plot(size: (4, 2),
     x-tick-step: none,
@@ -136,7 +134,7 @@
     padding: .1,
     stroke: black,
     fill: white,
-    orientation: ltr, default-position: "legend.north"))
+    orientation: ltr, default-position: "north"))
 
   plot.plot(size: (4, 2),
     x-tick-step: none,
@@ -164,3 +162,16 @@
       plot.add(samples: 10, ..dom, fn.with(offset: .2), mark: "|", style: (stroke: none), label: $ f_3(x) $)
     })
 })
+
+#test-case({
+  plot.plot(size: (4,2), x-tick-step: none, y-tick-step: none, {
+    plot.add(domain: (0,1), x => x)
+    plot.add-legend([Custom 1])
+    plot.add-legend([Custom 2], preview: () => {
+      import draw: *
+      set-style(stroke: blue)
+      line((0,0), (1,1))
+      line((0,1), (1,0))
+    })
+  })
+})