contribution_grid/
lib.rs

1//! A library for generating GitHub-style contribution graphs.
2//!
3//! This crate provides a builder interface to create contribution heatmaps, similar to those found on GitHub user profiles.
4//! It supports custom date ranges, colors, and dimensions, outputting the result as an image.
5//!
6//! # Examples
7//!
8//! ```rust
9//! use contribution_grid::{ContributionGraph, Theme, LinearStrategy};
10//! use chrono::NaiveDate;
11//! use std::collections::HashMap;
12//!
13//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
14//! let mut data = HashMap::new();
15//! data.insert(NaiveDate::from_ymd_opt(2023, 1, 1).unwrap(), 5);
16//! data.insert(NaiveDate::from_ymd_opt(2023, 1, 2).unwrap(), 12);
17//!
18//! // Use a built-in theme with a linear mapping strategy
19//! let img = ContributionGraph::new()
20//!     .with_data(data)
21//!     .theme(Theme::blue(LinearStrategy))
22//!     .generate();
23//!
24//! // img is an ImageBuffer from the `image` crate
25//! // img.save("graph.png")?;
26//! # Ok(())
27//! # }
28//! ```
29
30use std::collections::HashMap;
31
32use chrono::Datelike;
33use chrono::Duration;
34use chrono::NaiveDate;
35use image::ImageBuffer;
36use image::Rgba;
37use imageproc::drawing::draw_filled_rect_mut;
38use imageproc::rect::Rect;
39
40/// Defines how to map counts to color indices.
41pub trait MappingStrategy {
42    /// Maps a count to a color index `[0, num_colors)`.
43    fn map(&self, count: u32, max_count: u32, num_colors: usize) -> usize;
44}
45
46/// Linear percentile-based mapping.
47pub struct LinearStrategy;
48impl MappingStrategy for LinearStrategy {
49    fn map(&self, count: u32, max_count: u32, num_colors: usize) -> usize {
50        if count == 0 || max_count == 0 {
51            return 0;
52        }
53        let ratio = count as f32 / max_count as f32;
54        ((ratio * (num_colors - 1) as f32).round() as usize).min(num_colors - 1)
55    }
56}
57
58/// Logarithmic mapping (emphasizes differences at lower values).
59pub struct LogarithmicStrategy;
60impl MappingStrategy for LogarithmicStrategy {
61    fn map(&self, count: u32, max_count: u32, num_colors: usize) -> usize {
62        if count == 0 {
63            return 0;
64        }
65        let log_count = (count as f32 + 1.0).ln();
66        let log_max = (max_count as f32 + 1.0).ln();
67        let ratio = log_count / log_max;
68        ((ratio * (num_colors - 1) as f32).round() as usize).min(num_colors - 1)
69    }
70}
71
72/// Fixed threshold mapping (ignores `max_count`).
73pub struct ThresholdStrategy {
74    thresholds: Vec<u32>,
75}
76impl ThresholdStrategy {
77    /// Creates a new [`ThresholdStrategy`] with the given thresholds.
78    pub fn new(thresholds: Vec<u32>) -> Self {
79        Self { thresholds }
80    }
81}
82impl MappingStrategy for ThresholdStrategy {
83    fn map(&self, count: u32, _max_count: u32, _num_colors: usize) -> usize {
84        self.thresholds
85            .iter()
86            .position(|&t| count < t)
87            .unwrap_or(self.thresholds.len())
88    }
89}
90
91/// Defines the color palette and the strategy used to map counts to colors.
92pub struct Palette {
93    colors: Vec<Rgba<u8>>,
94    strategy: Box<dyn MappingStrategy>,
95}
96
97impl Palette {
98    /// Creates a new [`Palette`] with a list of colors and a mapping strategy.
99    pub fn new(colors: Vec<Rgba<u8>>, strategy: impl MappingStrategy + 'static) -> Self {
100        Self {
101            colors,
102            strategy: Box::new(strategy),
103        }
104    }
105
106    /// Returns the color for a given contribution count based on the global maximum count.
107    pub fn get_color(&self, count: u32, max_count: u32) -> Rgba<u8> {
108        let index = self.strategy.map(count, max_count, self.colors.len());
109        self.colors[index]
110    }
111}
112
113/// Helper for built-in color themes.
114pub struct Theme;
115
116impl Theme {
117    /// Modern GitHub green style (5 colors).
118    pub fn github(strategy: impl MappingStrategy + 'static) -> Palette {
119        Palette::new(
120            vec![
121                Rgba([235, 237, 240, 255]),
122                Rgba([155, 233, 168, 255]),
123                Rgba([64, 196, 99, 255]),
124                Rgba([48, 161, 78, 255]),
125                Rgba([33, 110, 57, 255]),
126            ],
127            strategy,
128        )
129    }
130
131    /// Classic GitHub style (5 colors).
132    pub fn github_old(strategy: impl MappingStrategy + 'static) -> Palette {
133        Palette::new(
134            vec![
135                Rgba([238, 238, 238, 255]),
136                Rgba([198, 228, 139, 255]),
137                Rgba([123, 201, 111, 255]),
138                Rgba([35, 154, 59, 255]),
139                Rgba([25, 97, 39, 255]),
140            ],
141            strategy,
142        )
143    }
144
145    /// A blue-shaded theme (7 colors).
146    pub fn blue(strategy: impl MappingStrategy + 'static) -> Palette {
147        Palette::new(
148            vec![
149                Rgba([235, 237, 240, 255]),
150                Rgba([201, 231, 245, 255]),
151                Rgba([168, 196, 238, 255]),
152                Rgba([135, 146, 232, 255]),
153                Rgba([96, 101, 181, 255]),
154                Rgba([61, 62, 125, 255]),
155                Rgba([31, 29, 66, 255]),
156            ],
157            strategy,
158        )
159    }
160
161    /// A red-shaded theme (5 colors).
162    pub fn red(strategy: impl MappingStrategy + 'static) -> Palette {
163        Palette::new(
164            vec![
165                Rgba([235, 237, 240, 255]),
166                Rgba([255, 208, 198, 255]),
167                Rgba([255, 148, 133, 255]),
168                Rgba([234, 86, 70, 255]),
169                Rgba([181, 28, 28, 255]),
170            ],
171            strategy,
172        )
173    }
174}
175
176/// A builder struct for generating GitHub-style contribution graphs.
177pub struct ContributionGraph {
178    data: HashMap<NaiveDate, u32>,
179    start_date: Option<NaiveDate>,
180    end_date: Option<NaiveDate>,
181    box_size: u32,
182    gap: u32,
183    margin: u32,
184    palette: Palette,
185    background_color: Rgba<u8>,
186    round_corners: bool,
187}
188
189impl Default for ContributionGraph {
190    fn default() -> Self {
191        Self {
192            data: HashMap::new(),
193            start_date: None,
194            end_date: None,
195            box_size: 11,
196            gap: 3,
197            margin: 20,
198            palette: Theme::github(LinearStrategy),
199            background_color: Rgba([0, 0, 0, 0]),
200            round_corners: true,
201        }
202    }
203}
204
205impl ContributionGraph {
206    /// Creates a new [`ContributionGraph`] builder with default settings.
207    pub fn new() -> Self {
208        Self::default()
209    }
210
211    /// Sets the contribution data for the graph.
212    ///
213    /// The data is a mapping from [`NaiveDate`] to a contribution count (`u32`).
214    pub fn with_data(mut self, data: HashMap<NaiveDate, u32>) -> Self {
215        self.data = data;
216        self
217    }
218
219    /// Sets the start date for the graph.
220    ///
221    /// If not provided, it defaults to the earliest date in the data or January 1st of the current year.
222    pub fn start_date(mut self, date: NaiveDate) -> Self {
223        self.start_date = Some(date);
224        self
225    }
226
227    /// Sets the end date for the graph.
228    ///
229    /// If not provided, it defaults to the latest date in the data or December 31st of the current year.
230    pub fn end_date(mut self, date: NaiveDate) -> Self {
231        self.end_date = Some(date);
232        self
233    }
234
235    /// Sets the size of each contribution box in pixels.
236    pub fn box_size(mut self, size: u32) -> Self {
237        self.box_size = size;
238        self
239    }
240
241    /// Sets the gap between contribution boxes in pixels.
242    pub fn gap(mut self, gap: u32) -> Self {
243        self.gap = gap;
244        self
245    }
246
247    /// Sets the margin around the entire grid in pixels.
248    pub fn margin(mut self, margin: u32) -> Self {
249        self.margin = margin;
250        self
251    }
252
253    /// Sets a predefined or custom [`Palette`].
254    ///
255    /// # Example
256    ///
257    /// ```rust
258    /// use contribution_grid::{ContributionGraph, Theme, LinearStrategy, Palette, ThresholdStrategy};
259    /// use image::Rgba;
260    ///
261    /// // Using a built-in theme
262    /// let graph = ContributionGraph::new()
263    ///     .theme(Theme::github(LinearStrategy));
264    ///
265    /// // Using a custom palette with thresholds
266    /// let custom_palette = Palette::new(
267    ///     vec![
268    ///         Rgba([20, 20, 20, 255]),
269    ///         Rgba([0, 255, 0, 255]),
270    ///         Rgba([0, 0, 255, 255]),
271    ///     ],
272    ///     ThresholdStrategy::new(vec![1, 10])
273    /// );
274    ///
275    /// let graph = ContributionGraph::new()
276    ///     .theme(custom_palette);
277    /// ```
278    pub fn theme(mut self, palette: Palette) -> Self {
279        self.palette = palette;
280        self
281    }
282
283    /// Sets the background color of the image.
284    pub fn background_color(mut self, color: Rgba<u8>) -> Self {
285        self.background_color = color;
286        self
287    }
288
289    /// Sets whether the contribution boxes should have rounded corners.
290    pub fn round_corners(mut self, round: bool) -> Self {
291        self.round_corners = round;
292        self
293    }
294
295    /// Calculates the effective start and end dates for the graph.
296    pub fn calculate_date_range(&self) -> (NaiveDate, NaiveDate) {
297        let current_year = chrono::Utc::now().naive_utc().year();
298        let start = self.start_date.unwrap_or_else(|| {
299            if self.data.is_empty() {
300                NaiveDate::from_ymd_opt(current_year, 1, 1).unwrap()
301            } else {
302                *self.data.keys().min().unwrap()
303            }
304        });
305
306        let end = self.end_date.unwrap_or_else(|| {
307            if self.data.is_empty() {
308                NaiveDate::from_ymd_opt(current_year, 12, 31).unwrap()
309            } else {
310                *self.data.keys().max().unwrap()
311            }
312        });
313
314        (start, end)
315    }
316
317    /// Calculates the dimensions (width, height) of the resulting image.
318    pub fn calculate_dimensions(&self, start: NaiveDate, end: NaiveDate) -> (u32, u32) {
319        let start_weekday = start.weekday().num_days_from_sunday() as i32;
320        let total_days = (end - start).num_days() as i32 + 1;
321        let weeks = (total_days + start_weekday + 6) / 7;
322
323        let width = self.margin * 2 + (weeks as u32) * (self.box_size + self.gap) - self.gap;
324        let height = self.margin * 2 + 7 * (self.box_size + self.gap) - self.gap;
325
326        (width, height)
327    }
328
329    /// Calculates the (x, y) coordinates for a given date relative to the start date.
330    pub fn get_coordinates(&self, date: NaiveDate, start: NaiveDate) -> (i32, i32) {
331        let start_weekday = start.weekday().num_days_from_sunday() as i32;
332        let day_idx = date.weekday().num_days_from_sunday() as i32;
333        let days_from_start = (date - start).num_days() as i32;
334        let week_idx = (days_from_start + start_weekday) / 7;
335
336        let x = self.margin as i32 + week_idx * (self.box_size as i32 + self.gap as i32);
337        let y = self.margin as i32 + day_idx * (self.box_size as i32 + self.gap as i32);
338
339        (x, y)
340    }
341
342    /// Draws a single contribution cell onto the image buffer.
343    fn draw_cell(&self, img: &mut ImageBuffer<Rgba<u8>, Vec<u8>>, x: i32, y: i32, color: Rgba<u8>) {
344        let width = img.width();
345        let height = img.height();
346
347        if x >= 0
348            && y >= 0
349            && (x as u32 + self.box_size) <= width
350            && (y as u32 + self.box_size) <= height
351        {
352            draw_filled_rect_mut(
353                img,
354                Rect::at(x, y).of_size(self.box_size, self.box_size),
355                color,
356            );
357
358            if self.round_corners && color[3] > 0 {
359                let corner_color = self.background_color;
360                if (x as u32) < width && (y as u32) < height {
361                    img.put_pixel(x as u32, y as u32, corner_color);
362                }
363                if (x as u32 + self.box_size - 1) < width && (y as u32) < height {
364                    img.put_pixel(
365                        (x + self.box_size as i32 - 1) as u32,
366                        y as u32,
367                        corner_color,
368                    );
369                }
370                if (x as u32) < width && (y as u32 + self.box_size - 1) < height {
371                    img.put_pixel(
372                        x as u32,
373                        (y + self.box_size as i32 - 1) as u32,
374                        corner_color,
375                    );
376                }
377                if (x as u32 + self.box_size - 1) < width && (y as u32 + self.box_size - 1) < height
378                {
379                    img.put_pixel(
380                        (x + self.box_size as i32 - 1) as u32,
381                        (y + self.box_size as i32 - 1) as u32,
382                        corner_color,
383                    );
384                }
385            }
386        }
387    }
388
389    /// Generates the contribution graph image based on the current configuration.
390    pub fn generate(&self) -> ImageBuffer<Rgba<u8>, Vec<u8>> {
391        let (start, end) = self.calculate_date_range();
392        let (width, height) = self.calculate_dimensions(start, end);
393
394        let max_count = self.data.values().cloned().max().unwrap_or(0);
395
396        let mut img = ImageBuffer::from_pixel(width, height, self.background_color);
397
398        let mut curr = start;
399        while curr <= end {
400            let count = self.data.get(&curr).cloned().unwrap_or(0);
401            let color = self.palette.get_color(count, max_count);
402
403            let (x, y) = self.get_coordinates(curr, start);
404            self.draw_cell(&mut img, x, y, color);
405
406            curr += Duration::days(1);
407        }
408
409        img
410    }
411}
412
413#[cfg(test)]
414mod tests {
415    use super::*;
416
417    #[test]
418    fn test_linear_mapping() {
419        let strategy = LinearStrategy;
420        // 0 counts -> index 0
421        assert_eq!(strategy.map(0, 100, 5), 0);
422        // 100 counts -> index 4
423        assert_eq!(strategy.map(100, 100, 5), 4);
424        // 50 counts -> index 2
425        assert_eq!(strategy.map(50, 100, 5), 2);
426    }
427
428    #[test]
429    fn test_threshold_mapping() {
430        let strategy = ThresholdStrategy::new(vec![1, 5, 10]);
431        assert_eq!(strategy.map(0, 0, 4), 0);
432        assert_eq!(strategy.map(3, 0, 4), 1);
433        assert_eq!(strategy.map(8, 0, 4), 2);
434        assert_eq!(strategy.map(15, 0, 4), 3);
435    }
436
437    #[test]
438    fn test_github_theme_colors() {
439        let palette = Theme::github(LinearStrategy);
440        assert_eq!(palette.get_color(0, 100), Rgba([235, 237, 240, 255]));
441        assert_eq!(palette.get_color(100, 100), Rgba([33, 110, 57, 255]));
442    }
443
444    #[test]
445    fn test_date_range_calculation() {
446        let graph = ContributionGraph::new();
447        let (start, end) = graph.calculate_date_range();
448        let year = chrono::Utc::now().naive_utc().year();
449        assert_eq!(start, NaiveDate::from_ymd_opt(year, 1, 1).unwrap());
450        assert_eq!(end, NaiveDate::from_ymd_opt(year, 12, 31).unwrap());
451    }
452
453    #[test]
454    fn test_dimensions_calculation() {
455        let start = NaiveDate::from_ymd_opt(2023, 1, 1).unwrap();
456        let end = NaiveDate::from_ymd_opt(2023, 1, 7).unwrap();
457        let graph = ContributionGraph::new().box_size(10).gap(2).margin(20);
458        let (width, height) = graph.calculate_dimensions(start, end);
459        assert_eq!(width, 50);
460        assert_eq!(height, 122);
461    }
462}
contribution_grid/lib.rs

contribution_grid/
lib.rs
