Skip to main content

unicode_plot/
plot.rs

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/// A labeled annotation attached to a row of the plot body.
10#[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    /// The annotation text.
19    #[must_use]
20    pub fn text(&self) -> &str {
21        &self.text
22    }
23
24    /// The annotation color, or `None` for the terminal's default foreground.
25    #[must_use]
26    pub const fn color(&self) -> Option<TermColor> {
27        self.color
28    }
29}
30
31/// Optional text placed at the six edge positions around the plot border
32/// (top-left, top-center, top-right, bottom-left, bottom-center, bottom-right).
33#[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/// All annotations attached to a plot: left-side labels, right-side labels,
77/// and edge decorations.
78#[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/// One of the six edge decoration positions around the plot border.
104#[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/// A complete plot: graphics area, annotations, border, title, and axis labels.
116///
117/// `Plot` owns its graphics and carries all metadata needed for rendering.
118/// Use the free-standing constructors ([`barplot`](crate::barplot()),
119/// [`lineplot`](crate::lineplot()), etc.) to create plots, then call
120/// [`render`](Plot::render) to write output.
121#[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    /// Default left margin (in characters) between the terminal edge and labels.
137    pub const DEFAULT_MARGIN: u16 = 3;
138    /// Default padding (in characters) between labels and the plot border.
139    pub const DEFAULT_PADDING: u16 = 1;
140
141    /// Creates a new plot wrapping the given graphics area with default settings.
142    #[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    /// Adds or replaces a labeled annotation on the left side of the plot body.
159    ///
160    /// Pass `None` for `color` to use the terminal's default foreground color.
161    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    /// Adds or replaces a labeled annotation on the right side of the plot body.
172    ///
173    /// Pass `None` for `color` to use the terminal's default foreground color.
174    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    /// Sets edge decoration text at one of the six border positions.
190    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    /// Returns the next color from the auto-series palette and advances the index.
203    #[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    /// Returns all annotations attached to this plot.
211    #[must_use]
212    pub const fn annotations(&self) -> &PlotAnnotations {
213        &self.annotations
214    }
215
216    /// Returns a reference to the underlying graphics area.
217    #[must_use]
218    pub fn graphics(&self) -> &G {
219        &self.graphics
220    }
221
222    /// Returns a mutable reference to the underlying graphics area.
223    pub fn graphics_mut(&mut self) -> &mut G {
224        &mut self.graphics
225    }
226
227    /// The plot title, rendered centered above the border.
228    #[must_use]
229    pub fn title(&self) -> Option<&str> {
230        self.title.as_deref()
231    }
232
233    /// The x-axis label, rendered centered below the border.
234    #[must_use]
235    pub fn xlabel(&self) -> Option<&str> {
236        self.xlabel.as_deref()
237    }
238
239    /// The y-axis label, rendered vertically along the left margin.
240    #[must_use]
241    pub fn ylabel(&self) -> Option<&str> {
242        self.ylabel.as_deref()
243    }
244
245    /// The border style for this plot.
246    #[must_use]
247    pub const fn border(&self) -> BorderType {
248        self.border
249    }
250
251    /// The left margin width in characters.
252    #[must_use]
253    pub const fn margin(&self) -> u16 {
254        self.margin
255    }
256
257    /// The padding between labels and the plot border.
258    #[must_use]
259    pub const fn padding(&self) -> u16 {
260        self.padding
261    }
262
263    /// Whether left/right row labels are rendered.
264    #[must_use]
265    pub const fn show_labels(&self) -> bool {
266        self.show_labels
267    }
268
269    /// The current auto-color palette index (wraps at 6).
270    #[must_use]
271    pub const fn auto_color_index(&self) -> u8 {
272        self.auto_color_index
273    }
274
275    /// Render this plot into the provided writer.
276    ///
277    /// # Errors
278    ///
279    /// Returns any I/O error produced by writing to `writer`.
280    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    /// Render this plot using a [`ColorMode`] policy.
296    ///
297    /// # Errors
298    ///
299    /// Returns any I/O error produced by writing to `writer`.
300    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    /// Render this plot using a [`ColorMode`] policy and terminal auto-detection.
316    ///
317    /// # Errors
318    ///
319    /// Returns any I/O error produced by writing to `writer`.
320    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}