1use std::collections::BTreeMap;
2use std::io::{IsTerminal, Result as IoResult, Write};
3
4use crate::border::BorderType;
5use crate::color::{AUTO_SERIES_COLORS, ColorMode, NamedColor, TermColor};
6use crate::graphics::GraphicsArea;
7use crate::render::{build_rendered_plot, write_ansi, write_plain};
8
9#[derive(Debug, Clone, PartialEq, Eq)]
11#[non_exhaustive]
12pub struct Annotation {
13 pub(crate) text: String,
14 pub(crate) color: Option<TermColor>,
15}
16
17impl Annotation {
18 #[must_use]
20 pub fn text(&self) -> &str {
21 &self.text
22 }
23
24 #[must_use]
26 pub const fn color(&self) -> Option<TermColor> {
27 self.color
28 }
29}
30
31#[derive(Debug, Clone, Default, PartialEq, Eq)]
34#[non_exhaustive]
35pub struct EdgeDecorations {
36 pub(crate) tl: Option<String>,
37 pub(crate) t: Option<String>,
38 pub(crate) tr: Option<String>,
39 pub(crate) bl: Option<String>,
40 pub(crate) b: Option<String>,
41 pub(crate) br: Option<String>,
42}
43
44impl EdgeDecorations {
45 #[must_use]
46 pub fn tl(&self) -> Option<&str> {
47 self.tl.as_deref()
48 }
49
50 #[must_use]
51 pub fn t(&self) -> Option<&str> {
52 self.t.as_deref()
53 }
54
55 #[must_use]
56 pub fn tr(&self) -> Option<&str> {
57 self.tr.as_deref()
58 }
59
60 #[must_use]
61 pub fn bl(&self) -> Option<&str> {
62 self.bl.as_deref()
63 }
64
65 #[must_use]
66 pub fn b(&self) -> Option<&str> {
67 self.b.as_deref()
68 }
69
70 #[must_use]
71 pub fn br(&self) -> Option<&str> {
72 self.br.as_deref()
73 }
74}
75
76#[derive(Debug, Clone, Default, PartialEq, Eq)]
79#[non_exhaustive]
80pub struct PlotAnnotations {
81 pub(crate) left: BTreeMap<usize, Annotation>,
82 pub(crate) right: BTreeMap<usize, Annotation>,
83 pub(crate) deco: EdgeDecorations,
84}
85
86impl PlotAnnotations {
87 #[must_use]
88 pub const fn left(&self) -> &BTreeMap<usize, Annotation> {
89 &self.left
90 }
91
92 #[must_use]
93 pub const fn right(&self) -> &BTreeMap<usize, Annotation> {
94 &self.right
95 }
96
97 #[must_use]
98 pub const fn decorations(&self) -> &EdgeDecorations {
99 &self.deco
100 }
101}
102
103#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
105#[non_exhaustive]
106pub enum DecorationPosition {
107 Tl,
108 T,
109 Tr,
110 Bl,
111 B,
112 Br,
113}
114
115#[non_exhaustive]
122pub struct Plot<G: GraphicsArea> {
123 pub(crate) graphics: G,
124 pub(crate) title: Option<String>,
125 pub(crate) xlabel: Option<String>,
126 pub(crate) ylabel: Option<String>,
127 pub(crate) border: BorderType,
128 pub(crate) margin: u16,
129 pub(crate) padding: u16,
130 pub(crate) show_labels: bool,
131 pub(crate) annotations: PlotAnnotations,
132 pub(crate) auto_color_index: u8,
133}
134
135impl<G: GraphicsArea> Plot<G> {
136 pub const DEFAULT_MARGIN: u16 = 3;
138 pub const DEFAULT_PADDING: u16 = 1;
140
141 #[must_use]
143 pub fn new(graphics: G) -> Self {
144 Self {
145 graphics,
146 title: None,
147 xlabel: None,
148 ylabel: None,
149 border: BorderType::Solid,
150 margin: Self::DEFAULT_MARGIN,
151 padding: Self::DEFAULT_PADDING,
152 show_labels: true,
153 annotations: PlotAnnotations::default(),
154 auto_color_index: 0,
155 }
156 }
157
158 pub fn annotate_left(&mut self, row: usize, text: impl Into<String>, color: Option<TermColor>) {
162 self.annotations.left.insert(
163 row,
164 Annotation {
165 text: text.into(),
166 color,
167 },
168 );
169 }
170
171 pub fn annotate_right(
175 &mut self,
176 row: usize,
177 text: impl Into<String>,
178 color: Option<TermColor>,
179 ) {
180 self.annotations.right.insert(
181 row,
182 Annotation {
183 text: text.into(),
184 color,
185 },
186 );
187 }
188
189 pub fn set_decoration(&mut self, position: DecorationPosition, text: impl Into<String>) {
191 let value = Some(text.into());
192 match position {
193 DecorationPosition::Tl => self.annotations.deco.tl = value,
194 DecorationPosition::T => self.annotations.deco.t = value,
195 DecorationPosition::Tr => self.annotations.deco.tr = value,
196 DecorationPosition::Bl => self.annotations.deco.bl = value,
197 DecorationPosition::B => self.annotations.deco.b = value,
198 DecorationPosition::Br => self.annotations.deco.br = value,
199 }
200 }
201
202 #[must_use]
204 pub fn next_color(&mut self) -> NamedColor {
205 let index = usize::from(self.auto_color_index) % AUTO_SERIES_COLORS.len();
206 self.auto_color_index = self.auto_color_index.wrapping_add(1);
207 AUTO_SERIES_COLORS[index]
208 }
209
210 #[must_use]
212 pub const fn annotations(&self) -> &PlotAnnotations {
213 &self.annotations
214 }
215
216 #[must_use]
218 pub fn graphics(&self) -> &G {
219 &self.graphics
220 }
221
222 pub fn graphics_mut(&mut self) -> &mut G {
224 &mut self.graphics
225 }
226
227 #[must_use]
229 pub fn title(&self) -> Option<&str> {
230 self.title.as_deref()
231 }
232
233 #[must_use]
235 pub fn xlabel(&self) -> Option<&str> {
236 self.xlabel.as_deref()
237 }
238
239 #[must_use]
241 pub fn ylabel(&self) -> Option<&str> {
242 self.ylabel.as_deref()
243 }
244
245 #[must_use]
247 pub const fn border(&self) -> BorderType {
248 self.border
249 }
250
251 #[must_use]
253 pub const fn margin(&self) -> u16 {
254 self.margin
255 }
256
257 #[must_use]
259 pub const fn padding(&self) -> u16 {
260 self.padding
261 }
262
263 #[must_use]
265 pub const fn show_labels(&self) -> bool {
266 self.show_labels
267 }
268
269 #[must_use]
271 pub const fn auto_color_index(&self) -> u8 {
272 self.auto_color_index
273 }
274
275 pub fn render(&self, writer: &mut impl Write, color: bool) -> IoResult<()> {
281 let color_mode = if color {
282 ColorMode::Always
283 } else {
284 ColorMode::Never
285 };
286 let use_color = should_use_color(color_mode, false);
287 let rendered = build_rendered_plot(self);
288 if use_color {
289 write_ansi(&rendered, writer)
290 } else {
291 write_plain(&rendered, writer)
292 }
293 }
294
295 pub fn render_with_mode(
301 &self,
302 writer: &mut impl Write,
303 color_mode: ColorMode,
304 writer_is_terminal: bool,
305 ) -> IoResult<()> {
306 let use_color = should_use_color(color_mode, writer_is_terminal);
307 let rendered = build_rendered_plot(self);
308 if use_color {
309 write_ansi(&rendered, writer)
310 } else {
311 write_plain(&rendered, writer)
312 }
313 }
314
315 pub fn render_with_mode_auto<W: Write + IsTerminal>(
321 &self,
322 writer: &mut W,
323 color_mode: ColorMode,
324 ) -> IoResult<()> {
325 self.render_with_mode(writer, color_mode, writer.is_terminal())
326 }
327}
328
329const fn should_use_color(color_mode: ColorMode, writer_is_terminal: bool) -> bool {
330 match color_mode {
331 ColorMode::Auto => writer_is_terminal,
332 ColorMode::Always => true,
333 ColorMode::Never => false,
334 }
335}
336
337#[cfg(test)]
338mod tests {
339 use super::{DecorationPosition, Plot, should_use_color};
340 use crate::color::{ColorMode, NamedColor, TermColor};
341 use crate::graphics::GraphicsArea;
342
343 #[derive(Debug, Default)]
344 struct DummyGraphics;
345
346 impl GraphicsArea for DummyGraphics {
347 fn nrows(&self) -> usize {
348 0
349 }
350
351 fn ncols(&self) -> usize {
352 0
353 }
354
355 fn render_row(&self, _row: usize, out: &mut crate::graphics::RowBuffer) {
356 out.clear();
357 }
358 }
359
360 #[test]
361 fn annotation_maps_are_deterministic_by_row() {
362 let mut plot = Plot::new(DummyGraphics);
363
364 plot.annotate_left(10, "left high", Some(TermColor::Named(NamedColor::Green)));
365 plot.annotate_left(2, "left low", Some(TermColor::Named(NamedColor::Blue)));
366 plot.annotate_right(7, "right mid", Some(TermColor::Named(NamedColor::Red)));
367 plot.annotate_right(1, "right low", Some(TermColor::Named(NamedColor::Cyan)));
368
369 let left_rows: Vec<_> = plot.annotations().left().keys().copied().collect();
370 let right_rows: Vec<_> = plot.annotations().right().keys().copied().collect();
371
372 assert_eq!(left_rows, vec![2, 10]);
373 assert_eq!(right_rows, vec![1, 7]);
374 }
375
376 #[test]
377 fn set_decoration_updates_expected_edge_slot() {
378 let mut plot = Plot::new(DummyGraphics);
379
380 plot.set_decoration(DecorationPosition::Tl, "max x");
381 plot.set_decoration(DecorationPosition::T, "top");
382 plot.set_decoration(DecorationPosition::Tr, "max y");
383 plot.set_decoration(DecorationPosition::Bl, "min x");
384 plot.set_decoration(DecorationPosition::B, "bottom");
385 plot.set_decoration(DecorationPosition::Br, "min y");
386
387 let deco = plot.annotations().decorations();
388 assert_eq!(deco.tl(), Some("max x"));
389 assert_eq!(deco.t(), Some("top"));
390 assert_eq!(deco.tr(), Some("max y"));
391 assert_eq!(deco.bl(), Some("min x"));
392 assert_eq!(deco.b(), Some("bottom"));
393 assert_eq!(deco.br(), Some("min y"));
394 }
395
396 #[test]
397 fn next_color_cycles_and_wraps_reference_sequence() {
398 let mut plot = Plot::new(DummyGraphics);
399 let sequence: Vec<_> = (0..8).map(|_| plot.next_color()).collect();
400
401 assert_eq!(
402 sequence,
403 vec![
404 NamedColor::Green,
405 NamedColor::Blue,
406 NamedColor::Red,
407 NamedColor::Magenta,
408 NamedColor::Yellow,
409 NamedColor::Cyan,
410 NamedColor::Green,
411 NamedColor::Blue,
412 ]
413 );
414 }
415
416 #[test]
417 fn annotate_left_replaces_existing_row_annotation() {
418 let mut plot = Plot::new(DummyGraphics);
419 plot.annotate_left(3, "first", Some(TermColor::Named(NamedColor::Green)));
420 plot.annotate_left(3, "second", Some(TermColor::Named(NamedColor::Red)));
421
422 let annotation = plot
423 .annotations()
424 .left()
425 .get(&3)
426 .unwrap_or_else(|| panic!("expected row 3 annotation"));
427 assert_eq!(annotation.text(), "second");
428 assert_eq!(annotation.color(), Some(TermColor::Named(NamedColor::Red)));
429 }
430
431 #[test]
432 fn next_color_handles_u8_overflow_and_palette_wrap() {
433 let mut plot = Plot::new(DummyGraphics);
434 plot.auto_color_index = u8::MAX;
435
436 assert_eq!(plot.next_color(), NamedColor::Magenta);
437 assert_eq!(plot.auto_color_index(), 0);
438 assert_eq!(plot.next_color(), NamedColor::Green);
439 assert_eq!(plot.auto_color_index(), 1);
440 }
441
442 #[test]
443 fn render_with_mode_auto_uses_writer_terminal_capability() {
444 assert!(should_use_color(ColorMode::Auto, true));
445 assert!(!should_use_color(ColorMode::Auto, false));
446 assert!(should_use_color(ColorMode::Always, false));
447 assert!(!should_use_color(ColorMode::Never, true));
448 }
449}