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//! # Core Types
7//!
8//! - [`MappingStrategy`] - Trait for mapping contribution counts to color indices
9//! - [`Palette`] - Defines colors and the strategy used to map counts to colors
10//! - [`ContributionGraph`] - Main builder for generating contribution graphs
11//!
12//! # Built-in Components
13//!
14//! The [`builtins`] module provides:
15//! - [`builtins::Theme`] - Factory for creating predefined color palettes
16//! - [`builtins::Strategy`] - Factory for creating built-in strategy instances
17//! - Built-in strategy implementations: [`builtins::LinearStrategy`], [`builtins::LogarithmicStrategy`], [`builtins::ThresholdStrategy`]
18//!
19//! # Examples
20//!
21//! ```rust
22//! use contribution_grid::{ContributionGraph, builtins::Theme, builtins::Strategy};
23//! use chrono::NaiveDate;
24//! use std::collections::HashMap;
25//!
26//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
27//! let mut data = HashMap::new();
28//! data.insert(NaiveDate::from_ymd_opt(2023, 1, 1).unwrap(), 5);
29//! data.insert(NaiveDate::from_ymd_opt(2023, 1, 2).unwrap(), 12);
30//!
31//! let img = ContributionGraph::new()
32//!     .with_data(data)
33//!     .theme(Theme::blue(Strategy::linear()))
34//!     .generate();
35//!
36//! // img.save("graph.png")?;
37//! # Ok(())
38//! # }
39//! ```
40
41use std::collections::HashMap;
42
43use chrono::Datelike;
44use chrono::Duration;
45use chrono::NaiveDate;
46use image::ImageBuffer;
47use image::Rgba;
48use imageproc::drawing::draw_filled_rect_mut;
49use imageproc::rect::Rect;
50
51pub mod builtins;
52
53/// Defines how to map counts to color indices.
54pub trait MappingStrategy {
55    /// Maps a count to a color index `[0, num_colors)`.
56    fn map(&self, count: u32, max_count: u32, num_colors: usize) -> usize;
57}
58
59/// Defines the color palette and the strategy used to map counts to colors.
60pub struct Palette {
61    colors: Vec<Rgba<u8>>,
62    strategy: Box<dyn MappingStrategy>,
63}
64
65impl Palette {
66    /// Creates a new [`Palette`] with a list of colors and a mapping strategy.
67    pub fn new(colors: Vec<Rgba<u8>>, strategy: impl MappingStrategy + 'static) -> Self {
68        Self {
69            colors,
70            strategy: Box::new(strategy),
71        }
72    }
73
74    /// Returns the color for a given contribution count based on the global maximum count.
75    pub fn get_color(&self, count: u32, max_count: u32) -> Rgba<u8> {
76        let index = self.strategy.map(count, max_count, self.colors.len());
77        self.colors[index]
78    }
79}
80
81/// A builder struct for generating GitHub-style contribution graphs.
82pub struct ContributionGraph {
83    data: HashMap<NaiveDate, u32>,
84    start_date: Option<NaiveDate>,
85    end_date: Option<NaiveDate>,
86    box_size: u32,
87    gap: u32,
88    margin: u32,
89    palette: Palette,
90    background_color: Rgba<u8>,
91    round_corners: bool,
92}
93
94impl Default for ContributionGraph {
95    fn default() -> Self {
96        Self {
97            data: HashMap::new(),
98            start_date: None,
99            end_date: None,
100            box_size: 11,
101            gap: 3,
102            margin: 20,
103            palette: builtins::Theme::github(builtins::Strategy::linear()),
104            background_color: Rgba([0, 0, 0, 0]),
105            round_corners: true,
106        }
107    }
108}
109
110impl ContributionGraph {
111    /// Creates a new [`ContributionGraph`] builder with default settings.
112    pub fn new() -> Self {
113        Self::default()
114    }
115
116    /// Sets the contribution data for the graph.
117    ///
118    /// The data is a mapping from [`NaiveDate`] to a contribution count (`u32`).
119    pub fn with_data(mut self, data: HashMap<NaiveDate, u32>) -> Self {
120        self.data = data;
121        self
122    }
123
124    /// Sets the start date for the graph.
125    ///
126    /// If not provided, it defaults to the earliest date in the data or January 1st of the current year.
127    pub fn start_date(mut self, date: NaiveDate) -> Self {
128        self.start_date = Some(date);
129        self
130    }
131
132    /// Sets the end date for the graph.
133    ///
134    /// If not provided, it defaults to the latest date in the data or December 31st of the current year.
135    pub fn end_date(mut self, date: NaiveDate) -> Self {
136        self.end_date = Some(date);
137        self
138    }
139
140    /// Sets the size of each contribution box in pixels.
141    pub fn box_size(mut self, size: u32) -> Self {
142        self.box_size = size;
143        self
144    }
145
146    /// Sets the gap between contribution boxes in pixels.
147    pub fn gap(mut self, gap: u32) -> Self {
148        self.gap = gap;
149        self
150    }
151
152    /// Sets the margin around the entire grid in pixels.
153    pub fn margin(mut self, margin: u32) -> Self {
154        self.margin = margin;
155        self
156    }
157
158    /// Sets a predefined or custom [`Palette`].
159    ///
160    /// # Example
161    ///
162    /// ```rust
163    /// use contribution_grid::{ContributionGraph, Palette};
164    /// use contribution_grid::builtins::{Theme, Strategy, ThresholdStrategy};
165    /// use image::Rgba;
166    ///
167    /// // Using a built-in theme
168    /// let graph = ContributionGraph::new()
169    ///     .theme(Theme::github(Strategy::linear()));
170    ///
171    /// // Using a custom palette with thresholds
172    /// let custom_palette = Palette::new(
173    ///     vec![
174    ///         Rgba([20, 20, 20, 255]),
175    ///         Rgba([0, 255, 0, 255]),
176    ///         Rgba([0, 0, 255, 255]),
177    ///     ],
178    ///     ThresholdStrategy::new(vec![1, 10])
179    /// );
180    ///
181    /// let graph = ContributionGraph::new()
182    ///     .theme(custom_palette);
183    /// ```
184    pub fn theme(mut self, palette: Palette) -> Self {
185        self.palette = palette;
186        self
187    }
188
189    /// Sets the background color of the image.
190    pub fn background_color(mut self, color: Rgba<u8>) -> Self {
191        self.background_color = color;
192        self
193    }
194
195    /// Sets whether the contribution boxes should have rounded corners.
196    pub fn round_corners(mut self, round: bool) -> Self {
197        self.round_corners = round;
198        self
199    }
200
201    /// Generates the contribution graph image based on the current configuration.
202    pub fn generate(&self) -> ImageBuffer<Rgba<u8>, Vec<u8>> {
203        let (start, end) = self.calculate_date_range();
204        let (width, height) = self.calculate_dimensions(start, end);
205
206        let max_count = self.data.values().cloned().max().unwrap_or(0);
207
208        let mut img = ImageBuffer::from_pixel(width, height, self.background_color);
209
210        let mut curr = start;
211        while curr <= end {
212            let count = self.data.get(&curr).cloned().unwrap_or(0);
213            let color = self.palette.get_color(count, max_count);
214
215            let (x, y) = self.get_coordinates(curr, start);
216            self.draw_cell(&mut img, x, y, color);
217
218            curr += Duration::days(1);
219        }
220
221        img
222    }
223
224    /// Calculates the effective start and end dates for the graph.
225    fn calculate_date_range(&self) -> (NaiveDate, NaiveDate) {
226        let current_year = chrono::Utc::now().naive_utc().year();
227        let start = self.start_date.unwrap_or_else(|| {
228            if self.data.is_empty() {
229                NaiveDate::from_ymd_opt(current_year, 1, 1).unwrap()
230            } else {
231                *self.data.keys().min().unwrap()
232            }
233        });
234
235        let end = self.end_date.unwrap_or_else(|| {
236            if self.data.is_empty() {
237                NaiveDate::from_ymd_opt(current_year, 12, 31).unwrap()
238            } else {
239                *self.data.keys().max().unwrap()
240            }
241        });
242
243        (start, end)
244    }
245
246    /// Calculates the dimensions (width, height) of the resulting image.
247    fn calculate_dimensions(&self, start: NaiveDate, end: NaiveDate) -> (u32, u32) {
248        let start_weekday = start.weekday().num_days_from_sunday() as i32;
249        let total_days = (end - start).num_days() as i32 + 1;
250        let weeks = (total_days + start_weekday + 6) / 7;
251
252        let width = self.margin * 2 + (weeks as u32) * (self.box_size + self.gap) - self.gap;
253        let height = self.margin * 2 + 7 * (self.box_size + self.gap) - self.gap;
254
255        (width, height)
256    }
257
258    /// Calculates the (x, y) coordinates for a given date relative to the start date.
259    fn get_coordinates(&self, date: NaiveDate, start: NaiveDate) -> (i32, i32) {
260        let start_weekday = start.weekday().num_days_from_sunday() as i32;
261        let day_idx = date.weekday().num_days_from_sunday() as i32;
262        let days_from_start = (date - start).num_days() as i32;
263        let week_idx = (days_from_start + start_weekday) / 7;
264
265        let x = self.margin as i32 + week_idx * (self.box_size as i32 + self.gap as i32);
266        let y = self.margin as i32 + day_idx * (self.box_size as i32 + self.gap as i32);
267
268        (x, y)
269    }
270
271    /// Draws a single contribution cell onto the image buffer.
272    fn draw_cell(&self, img: &mut ImageBuffer<Rgba<u8>, Vec<u8>>, x: i32, y: i32, color: Rgba<u8>) {
273        let width = img.width();
274        let height = img.height();
275
276        if x >= 0
277            && y >= 0
278            && (x as u32 + self.box_size) <= width
279            && (y as u32 + self.box_size) <= height
280        {
281            draw_filled_rect_mut(
282                img,
283                Rect::at(x, y).of_size(self.box_size, self.box_size),
284                color,
285            );
286
287            if self.round_corners && color[3] > 0 {
288                let corner_color = self.background_color;
289                if (x as u32) < width && (y as u32) < height {
290                    img.put_pixel(x as u32, y as u32, corner_color);
291                }
292                if (x as u32 + self.box_size - 1) < width && (y as u32) < height {
293                    img.put_pixel(
294                        (x + self.box_size as i32 - 1) as u32,
295                        y as u32,
296                        corner_color,
297                    );
298                }
299                if (x as u32) < width && (y as u32 + self.box_size - 1) < height {
300                    img.put_pixel(
301                        x as u32,
302                        (y + self.box_size as i32 - 1) as u32,
303                        corner_color,
304                    );
305                }
306                if (x as u32 + self.box_size - 1) < width && (y as u32 + self.box_size - 1) < height
307                {
308                    img.put_pixel(
309                        (x + self.box_size as i32 - 1) as u32,
310                        (y + self.box_size as i32 - 1) as u32,
311                        corner_color,
312                    );
313                }
314            }
315        }
316    }
317}
318
319#[cfg(test)]
320mod tests {
321    use super::*;
322
323    #[test]
324    fn test_linear_mapping() {
325        let strategy = builtins::LinearStrategy;
326        // 0 counts -> index 0
327        assert_eq!(strategy.map(0, 100, 5), 0);
328        // 100 counts -> index 4
329        assert_eq!(strategy.map(100, 100, 5), 4);
330        // 50 counts -> index 2
331        assert_eq!(strategy.map(50, 100, 5), 2);
332    }
333
334    #[test]
335    fn test_threshold_mapping() {
336        let strategy = builtins::ThresholdStrategy::new(vec![1, 5, 10]);
337        assert_eq!(strategy.map(0, 0, 4), 0);
338        assert_eq!(strategy.map(3, 0, 4), 1);
339        assert_eq!(strategy.map(8, 0, 4), 2);
340        assert_eq!(strategy.map(15, 0, 4), 3);
341    }
342
343    #[test]
344    fn test_github_theme_colors() {
345        let palette = builtins::Theme::github(builtins::Strategy::linear());
346        assert_eq!(palette.get_color(0, 100), Rgba([235, 237, 240, 255]));
347        assert_eq!(palette.get_color(100, 100), Rgba([33, 110, 57, 255]));
348    }
349
350    #[test]
351    fn test_date_range_calculation() {
352        let graph = ContributionGraph::new();
353        let (start, end) = graph.calculate_date_range();
354        let year = chrono::Utc::now().naive_utc().year();
355        assert_eq!(start, NaiveDate::from_ymd_opt(year, 1, 1).unwrap());
356        assert_eq!(end, NaiveDate::from_ymd_opt(year, 12, 31).unwrap());
357    }
358
359    #[test]
360    fn test_dimensions_calculation() {
361        let start = NaiveDate::from_ymd_opt(2023, 1, 1).unwrap();
362        let end = NaiveDate::from_ymd_opt(2023, 1, 7).unwrap();
363        let graph = ContributionGraph::new().box_size(10).gap(2).margin(20);
364        let (width, height) = graph.calculate_dimensions(start, end);
365        assert_eq!(width, 50);
366        assert_eq!(height, 122);
367    }
368}
contribution_grid/lib.rs

contribution_grid/
lib.rs
