Skip to main content

rumatui_tui/widgets/
chart.rs

1use crate::{
2    buffer::Buffer,
3    layout::{Constraint, Rect},
4    style::Style,
5    symbols,
6    widgets::{
7        canvas::{Canvas, Line, Points},
8        Block, Borders, Widget,
9    },
10};
11use std::{borrow::Cow, cmp::max};
12use unicode_width::UnicodeWidthStr;
13
14/// An X or Y axis for the chart widget
15pub struct Axis<'a, L>
16where
17    L: AsRef<str> + 'a,
18{
19    /// Title displayed next to axis end
20    title: Option<&'a str>,
21    /// Style of the title
22    title_style: Style,
23    /// Bounds for the axis (all data points outside these limits will not be represented)
24    bounds: [f64; 2],
25    /// A list of labels to put to the left or below the axis
26    labels: Option<&'a [L]>,
27    /// The labels' style
28    labels_style: Style,
29    /// The style used to draw the axis itself
30    style: Style,
31}
32
33impl<'a, L> Default for Axis<'a, L>
34where
35    L: AsRef<str>,
36{
37    fn default() -> Axis<'a, L> {
38        Axis {
39            title: None,
40            title_style: Default::default(),
41            bounds: [0.0, 0.0],
42            labels: None,
43            labels_style: Default::default(),
44            style: Default::default(),
45        }
46    }
47}
48
49impl<'a, L> Axis<'a, L>
50where
51    L: AsRef<str>,
52{
53    pub fn title(mut self, title: &'a str) -> Axis<'a, L> {
54        self.title = Some(title);
55        self
56    }
57
58    pub fn title_style(mut self, style: Style) -> Axis<'a, L> {
59        self.title_style = style;
60        self
61    }
62
63    pub fn bounds(mut self, bounds: [f64; 2]) -> Axis<'a, L> {
64        self.bounds = bounds;
65        self
66    }
67
68    pub fn labels(mut self, labels: &'a [L]) -> Axis<'a, L> {
69        self.labels = Some(labels);
70        self
71    }
72
73    pub fn labels_style(mut self, style: Style) -> Axis<'a, L> {
74        self.labels_style = style;
75        self
76    }
77
78    pub fn style(mut self, style: Style) -> Axis<'a, L> {
79        self.style = style;
80        self
81    }
82}
83
84/// Marker to use when plotting data points
85pub enum Marker {
86    /// One point per cell
87    Dot,
88    /// Up to 8 points per cell
89    Braille,
90}
91
92/// Used to determine which style of graphing to use
93pub enum GraphType {
94    /// Draw each point
95    Scatter,
96    /// Draw each point and lines between each point using the same marker
97    Line,
98}
99
100/// A group of data points
101pub struct Dataset<'a> {
102    /// Name of the dataset (used in the legend if shown)
103    name: Cow<'a, str>,
104    /// A reference to the actual data
105    data: &'a [(f64, f64)],
106    /// Symbol used for each points of this dataset
107    marker: Marker,
108    /// Determines graph type used for drawing points
109    graph_type: GraphType,
110    /// Style used to plot this dataset
111    style: Style,
112}
113
114impl<'a> Default for Dataset<'a> {
115    fn default() -> Dataset<'a> {
116        Dataset {
117            name: Cow::from(""),
118            data: &[],
119            marker: Marker::Dot,
120            graph_type: GraphType::Scatter,
121            style: Style::default(),
122        }
123    }
124}
125
126impl<'a> Dataset<'a> {
127    pub fn name<S>(mut self, name: S) -> Dataset<'a>
128    where
129        S: Into<Cow<'a, str>>,
130    {
131        self.name = name.into();
132        self
133    }
134
135    pub fn data(mut self, data: &'a [(f64, f64)]) -> Dataset<'a> {
136        self.data = data;
137        self
138    }
139
140    pub fn marker(mut self, marker: Marker) -> Dataset<'a> {
141        self.marker = marker;
142        self
143    }
144
145    pub fn graph_type(mut self, graph_type: GraphType) -> Dataset<'a> {
146        self.graph_type = graph_type;
147        self
148    }
149
150    pub fn style(mut self, style: Style) -> Dataset<'a> {
151        self.style = style;
152        self
153    }
154}
155
156/// A container that holds all the infos about where to display each elements of the chart (axis,
157/// labels, legend, ...).
158#[derive(Debug, Clone, PartialEq)]
159struct ChartLayout {
160    /// Location of the title of the x axis
161    title_x: Option<(u16, u16)>,
162    /// Location of the title of the y axis
163    title_y: Option<(u16, u16)>,
164    /// Location of the first label of the x axis
165    label_x: Option<u16>,
166    /// Location of the first label of the y axis
167    label_y: Option<u16>,
168    /// Y coordinate of the horizontal axis
169    axis_x: Option<u16>,
170    /// X coordinate of the vertical axis
171    axis_y: Option<u16>,
172    /// Area of the legend
173    legend_area: Option<Rect>,
174    /// Area of the graph
175    graph_area: Rect,
176}
177
178impl Default for ChartLayout {
179    fn default() -> ChartLayout {
180        ChartLayout {
181            title_x: None,
182            title_y: None,
183            label_x: None,
184            label_y: None,
185            axis_x: None,
186            axis_y: None,
187            legend_area: None,
188            graph_area: Rect::default(),
189        }
190    }
191}
192
193/// A widget to plot one or more dataset in a cartesian coordinate system
194///
195/// # Examples
196///
197/// ```
198/// # use rumatui_tui::widgets::{Block, Borders, Chart, Axis, Dataset, Marker, GraphType};
199/// # use rumatui_tui::style::{Style, Color};
200/// Chart::default()
201///     .block(Block::default().title("Chart"))
202///     .x_axis(Axis::default()
203///         .title("X Axis")
204///         .title_style(Style::default().fg(Color::Red))
205///         .style(Style::default().fg(Color::White))
206///         .bounds([0.0, 10.0])
207///         .labels(&["0.0", "5.0", "10.0"]))
208///     .y_axis(Axis::default()
209///         .title("Y Axis")
210///         .title_style(Style::default().fg(Color::Red))
211///         .style(Style::default().fg(Color::White))
212///         .bounds([0.0, 10.0])
213///         .labels(&["0.0", "5.0", "10.0"]))
214///     .datasets(&[Dataset::default()
215///                     .name("data1")
216///                     .marker(Marker::Dot)
217///                     .graph_type(GraphType::Scatter)
218///                     .style(Style::default().fg(Color::Cyan))
219///                     .data(&[(0.0, 5.0), (1.0, 6.0), (1.5, 6.434)]),
220///                 Dataset::default()
221///                     .name("data2")
222///                     .marker(Marker::Braille)
223///                     .graph_type(GraphType::Line)
224///                     .style(Style::default().fg(Color::Magenta))
225///                     .data(&[(4.0, 5.0), (5.0, 8.0), (7.66, 13.5)])]);
226/// ```
227pub struct Chart<'a, LX, LY>
228where
229    LX: AsRef<str> + 'a,
230    LY: AsRef<str> + 'a,
231{
232    /// A block to display around the widget eventually
233    block: Option<Block<'a>>,
234    /// The horizontal axis
235    x_axis: Axis<'a, LX>,
236    /// The vertical axis
237    y_axis: Axis<'a, LY>,
238    /// A reference to the datasets
239    datasets: &'a [Dataset<'a>],
240    /// The widget base style
241    style: Style,
242    /// Constraints used to determine whether the legend should be shown or
243    /// not
244    hidden_legend_constraints: (Constraint, Constraint),
245}
246
247impl<'a, LX, LY> Default for Chart<'a, LX, LY>
248where
249    LX: AsRef<str>,
250    LY: AsRef<str>,
251{
252    fn default() -> Chart<'a, LX, LY> {
253        Chart {
254            block: None,
255            x_axis: Axis::default(),
256            y_axis: Axis::default(),
257            style: Default::default(),
258            datasets: &[],
259            hidden_legend_constraints: (Constraint::Ratio(1, 4), Constraint::Ratio(1, 4)),
260        }
261    }
262}
263
264impl<'a, LX, LY> Chart<'a, LX, LY>
265where
266    LX: AsRef<str>,
267    LY: AsRef<str>,
268{
269    pub fn block(mut self, block: Block<'a>) -> Chart<'a, LX, LY> {
270        self.block = Some(block);
271        self
272    }
273
274    pub fn style(mut self, style: Style) -> Chart<'a, LX, LY> {
275        self.style = style;
276        self
277    }
278
279    pub fn x_axis(mut self, axis: Axis<'a, LX>) -> Chart<'a, LX, LY> {
280        self.x_axis = axis;
281        self
282    }
283
284    pub fn y_axis(mut self, axis: Axis<'a, LY>) -> Chart<'a, LX, LY> {
285        self.y_axis = axis;
286        self
287    }
288
289    pub fn datasets(mut self, datasets: &'a [Dataset<'a>]) -> Chart<'a, LX, LY> {
290        self.datasets = datasets;
291        self
292    }
293
294    /// Set the constraints used to determine whether the legend should be shown or
295    /// not.
296    ///
297    /// # Examples
298    ///
299    /// ```
300    /// # use rumatui_tui::widgets::Chart;
301    /// # use rumatui_tui::layout::Constraint;
302    /// let constraints = (
303    ///     Constraint::Ratio(1, 3),
304    ///     Constraint::Ratio(1, 4)
305    /// );
306    /// // Hide the legend when either its width is greater than 33% of the total widget width
307    /// // or if its height is greater than 25% of the total widget height.
308    /// let _chart: Chart<String, String> = Chart::default()
309    ///     .hidden_legend_constraints(constraints);
310    pub fn hidden_legend_constraints(
311        mut self,
312        constraints: (Constraint, Constraint),
313    ) -> Chart<'a, LX, LY> {
314        self.hidden_legend_constraints = constraints;
315        self
316    }
317
318    /// Compute the internal layout of the chart given the area. If the area is too small some
319    /// elements may be automatically hidden
320    fn layout(&self, area: Rect) -> ChartLayout {
321        let mut layout = ChartLayout::default();
322        if area.height == 0 || area.width == 0 {
323            return layout;
324        }
325        let mut x = area.left();
326        let mut y = area.bottom() - 1;
327
328        if self.x_axis.labels.is_some() && y > area.top() {
329            layout.label_x = Some(y);
330            y -= 1;
331        }
332
333        if let Some(y_labels) = self.y_axis.labels {
334            let mut max_width = y_labels
335                .iter()
336                .fold(0, |acc, l| max(l.as_ref().width(), acc))
337                as u16;
338            if let Some(x_labels) = self.x_axis.labels {
339                if !x_labels.is_empty() {
340                    max_width = max(max_width, x_labels[0].as_ref().width() as u16);
341                }
342            }
343            if x + max_width < area.right() {
344                layout.label_y = Some(x);
345                x += max_width;
346            }
347        }
348
349        if self.x_axis.labels.is_some() && y > area.top() {
350            layout.axis_x = Some(y);
351            y -= 1;
352        }
353
354        if self.y_axis.labels.is_some() && x + 1 < area.right() {
355            layout.axis_y = Some(x);
356            x += 1;
357        }
358
359        if x < area.right() && y > 1 {
360            layout.graph_area = Rect::new(x, area.top(), area.right() - x, y - area.top() + 1);
361        }
362
363        if let Some(title) = self.x_axis.title {
364            let w = title.width() as u16;
365            if w < layout.graph_area.width && layout.graph_area.height > 2 {
366                layout.title_x = Some((x + layout.graph_area.width - w, y));
367            }
368        }
369
370        if let Some(title) = self.y_axis.title {
371            let w = title.width() as u16;
372            if w + 1 < layout.graph_area.width && layout.graph_area.height > 2 {
373                layout.title_y = Some((x + 1, area.top()));
374            }
375        }
376
377        if let Some(inner_width) = self.datasets.iter().map(|d| d.name.width() as u16).max() {
378            let legend_width = inner_width + 2;
379            let legend_height = self.datasets.len() as u16 + 2;
380            let max_legend_width = self
381                .hidden_legend_constraints
382                .0
383                .apply(layout.graph_area.width);
384            let max_legend_height = self
385                .hidden_legend_constraints
386                .1
387                .apply(layout.graph_area.height);
388            if inner_width > 0
389                && legend_width < max_legend_width
390                && legend_height < max_legend_height
391            {
392                layout.legend_area = Some(Rect::new(
393                    layout.graph_area.right() - legend_width,
394                    layout.graph_area.top(),
395                    legend_width,
396                    legend_height,
397                ));
398            }
399        }
400        layout
401    }
402}
403
404impl<'a, LX, LY> Widget for Chart<'a, LX, LY>
405where
406    LX: AsRef<str>,
407    LY: AsRef<str>,
408{
409    fn render(mut self, area: Rect, buf: &mut Buffer) {
410        let chart_area = match self.block {
411            Some(ref mut b) => {
412                b.render(area, buf);
413                b.inner(area)
414            }
415            None => area,
416        };
417
418        let layout = self.layout(chart_area);
419        let graph_area = layout.graph_area;
420        if graph_area.width < 1 || graph_area.height < 1 {
421            return;
422        }
423
424        buf.set_background(chart_area, self.style.bg);
425
426        if let Some((x, y)) = layout.title_x {
427            let title = self.x_axis.title.unwrap();
428            buf.set_string(x, y, title, self.x_axis.title_style);
429        }
430
431        if let Some((x, y)) = layout.title_y {
432            let title = self.y_axis.title.unwrap();
433            buf.set_string(x, y, title, self.y_axis.title_style);
434        }
435
436        if let Some(y) = layout.label_x {
437            let labels = self.x_axis.labels.unwrap();
438            let total_width = labels.iter().fold(0, |acc, l| l.as_ref().width() + acc) as u16;
439            let labels_len = labels.len() as u16;
440            if total_width < graph_area.width && labels_len > 1 {
441                for (i, label) in labels.iter().enumerate() {
442                    buf.set_string(
443                        graph_area.left() + i as u16 * (graph_area.width - 1) / (labels_len - 1)
444                            - label.as_ref().width() as u16,
445                        y,
446                        label.as_ref(),
447                        self.x_axis.labels_style,
448                    );
449                }
450            }
451        }
452
453        if let Some(x) = layout.label_y {
454            let labels = self.y_axis.labels.unwrap();
455            let labels_len = labels.len() as u16;
456            for (i, label) in labels.iter().enumerate() {
457                let dy = i as u16 * (graph_area.height - 1) / (labels_len - 1);
458                if dy < graph_area.bottom() {
459                    buf.set_string(
460                        x,
461                        graph_area.bottom() - 1 - dy,
462                        label.as_ref(),
463                        self.y_axis.labels_style,
464                    );
465                }
466            }
467        }
468
469        if let Some(y) = layout.axis_x {
470            for x in graph_area.left()..graph_area.right() {
471                buf.get_mut(x, y)
472                    .set_symbol(symbols::line::HORIZONTAL)
473                    .set_style(self.x_axis.style);
474            }
475        }
476
477        if let Some(x) = layout.axis_y {
478            for y in graph_area.top()..graph_area.bottom() {
479                buf.get_mut(x, y)
480                    .set_symbol(symbols::line::VERTICAL)
481                    .set_style(self.y_axis.style);
482            }
483        }
484
485        if let Some(y) = layout.axis_x {
486            if let Some(x) = layout.axis_y {
487                buf.get_mut(x, y)
488                    .set_symbol(symbols::line::BOTTOM_LEFT)
489                    .set_style(self.x_axis.style);
490            }
491        }
492
493        for dataset in self.datasets {
494            match dataset.marker {
495                Marker::Dot => {
496                    for &(x, y) in dataset.data.iter().filter(|&&(x, y)| {
497                        !(x < self.x_axis.bounds[0]
498                            || x > self.x_axis.bounds[1]
499                            || y < self.y_axis.bounds[0]
500                            || y > self.y_axis.bounds[1])
501                    }) {
502                        let dy = ((self.y_axis.bounds[1] - y) * f64::from(graph_area.height - 1)
503                            / (self.y_axis.bounds[1] - self.y_axis.bounds[0]))
504                            as u16;
505                        let dx = ((x - self.x_axis.bounds[0]) * f64::from(graph_area.width - 1)
506                            / (self.x_axis.bounds[1] - self.x_axis.bounds[0]))
507                            as u16;
508
509                        buf.get_mut(graph_area.left() + dx, graph_area.top() + dy)
510                            .set_symbol(symbols::DOT)
511                            .set_fg(dataset.style.fg)
512                            .set_bg(dataset.style.bg);
513                    }
514                }
515                Marker::Braille => {
516                    Canvas::default()
517                        .background_color(self.style.bg)
518                        .x_bounds(self.x_axis.bounds)
519                        .y_bounds(self.y_axis.bounds)
520                        .paint(|ctx| {
521                            ctx.draw(&Points {
522                                coords: dataset.data,
523                                color: dataset.style.fg,
524                            });
525                            if let GraphType::Line = dataset.graph_type {
526                                for i in 0..dataset.data.len() - 1 {
527                                    ctx.draw(&Line {
528                                        x1: dataset.data[i].0,
529                                        y1: dataset.data[i].1,
530                                        x2: dataset.data[i + 1].0,
531                                        y2: dataset.data[i + 1].1,
532                                        color: dataset.style.fg,
533                                    })
534                                }
535                            }
536                        })
537                        .render(graph_area, buf);
538                }
539            }
540        }
541
542        if let Some(legend_area) = layout.legend_area {
543            Block::default()
544                .borders(Borders::ALL)
545                .render(legend_area, buf);
546            for (i, dataset) in self.datasets.iter().enumerate() {
547                buf.set_string(
548                    legend_area.x + 1,
549                    legend_area.y + 1 + i as u16,
550                    &dataset.name,
551                    dataset.style,
552                );
553            }
554        }
555    }
556}
557
558#[cfg(test)]
559mod tests {
560    use super::*;
561
562    struct LegendTestCase {
563        chart_area: Rect,
564        hidden_legend_constraints: (Constraint, Constraint),
565        legend_area: Option<Rect>,
566    }
567
568    #[test]
569    fn it_should_hide_the_legend() {
570        let data = [(0.0, 5.0), (1.0, 6.0), (3.0, 7.0)];
571        let datasets = (0..10)
572            .map(|i| {
573                let name = format!("Dataset #{}", i);
574                Dataset::default().name(name).data(&data)
575            })
576            .collect::<Vec<_>>();
577        let cases = [
578            LegendTestCase {
579                chart_area: Rect::new(0, 0, 100, 100),
580                hidden_legend_constraints: (Constraint::Ratio(1, 4), Constraint::Ratio(1, 4)),
581                legend_area: Some(Rect::new(88, 0, 12, 12)),
582            },
583            LegendTestCase {
584                chart_area: Rect::new(0, 0, 100, 100),
585                hidden_legend_constraints: (Constraint::Ratio(1, 10), Constraint::Ratio(1, 4)),
586                legend_area: None,
587            },
588        ];
589        for case in &cases {
590            let chart: Chart<String, String> = Chart::default()
591                .x_axis(Axis::default().title("X axis"))
592                .y_axis(Axis::default().title("Y axis"))
593                .hidden_legend_constraints(case.hidden_legend_constraints)
594                .datasets(datasets.as_slice());
595            let layout = chart.layout(case.chart_area);
596            assert_eq!(layout.legend_area, case.legend_area);
597        }
598    }
599}