Chunk based heatmap for wrapping on multiple rows

src/cli.rs

@@ -38,6 +38,13 @@ pub struct CliArgs {
 	#[arg(long("split-months"), help("Split months"), default_value_t = false)]
 	pub split_months: bool,
 
+	#[arg(
+		long("months-per-row"),
+		help("Wrap the months on a new row"),
+		default_value_t = 13
+	)]
+	pub months_per_row: u16,
+
 	#[arg(long("no-merges"), default_value_t = false)]
 	pub no_merges: bool,
 

src/heatmap.rs

@@ -1,4 +1,7 @@
-use std::{collections::BTreeMap, fmt::Display};
+use std::{
+	collections::BTreeMap,
+	fmt::{Display, Write},
+};
 
 use chrono::{Datelike, Duration, NaiveDate};
 use clap::ValueEnum;
@@ -7,13 +10,12 @@ use itertools::Itertools;
 use crate::{get_char, get_color, get_color_map, Commit, DAYS, RESET};
 
 pub struct Heatmap {
-	data: [Vec<i32>; 7],
 	since: NaiveDate,
 	until: NaiveDate,
 	commits: Vec<Commit>,
-	months: Vec<(usize, String)>,
 	highest_count: i32,
 	repos: usize,
+	chunks: Vec<Chunk>,
 
 	format: Format,
 }
@@ -25,16 +27,16 @@ impl Heatmap {
 		repos: usize,
 		commits: Vec<Commit>,
 		split_months: bool,
+		months_per_row: u16,
 		format: Format,
 	) -> Self {
 		let mut heatmap = Self {
-			data: [vec![], vec![], vec![], vec![], vec![], vec![], vec![]],
 			since,
 			until,
 			commits,
-			months: vec![],
 			highest_count: 0,
 			repos,
+			chunks: vec![],
 			format,
 		};
 
@@ -49,11 +51,8 @@ impl Heatmap {
 		let mut current_day = since;
 		let mut day_of_week = current_day.weekday().num_days_from_monday() % 7;
 
-		if day_of_week != 0 {
-			for i in 0..day_of_week {
-				heatmap.data[i as usize].push(-1);
-			}
-		}
+		let mut chunk = Chunk::new(day_of_week);
+		let mut chunk_idx = 0;
 
 		// Track the very first day of the heatmap, as we don't want the extra spacing in front of
 		// those.
@@ -65,7 +64,7 @@ impl Heatmap {
 				// we add 2 weeks worth of empty space so months are more visible
 				if !first_day && current_day.day0() == 0 {
 					for i in 0..14 {
-						heatmap.data[(i as usize) % 7].push(-1);
+						chunk.data[(i as usize) % 7].push(-1);
 					}
 				}
 				first_day = false;
@@ -74,53 +73,44 @@ impl Heatmap {
 			let month_name = current_day.format("%b").to_string();
 
 			if current_day == since {
-				heatmap.months.push((0, month_name));
+				chunk.months.push((0, month_name));
 			}
 			else if current_day.day0() == 0 {
-				heatmap
+				chunk_idx += 1;
+
+				if chunk_idx > months_per_row - 1 {
+					heatmap.chunks.push(chunk);
+					chunk = Chunk::new(day_of_week);
+					chunk_idx = 0;
+				}
+
+				chunk
 					.months
-					.push((heatmap.data[day_of_week as usize].len(), month_name));
+					.push((chunk.data[day_of_week as usize].len(), month_name));
 			}
 
 			let value = grouped_commits.get(&current_day);
 			match value {
 				Some(val) => {
-					heatmap.data[day_of_week as usize].push(*val);
+					chunk.data[day_of_week as usize].push(*val);
 					if *val > heatmap.highest_count {
 						heatmap.highest_count = *val;
 					}
 				}
-				None => heatmap.data[day_of_week as usize].push(0),
+				None => {
+					chunk.data[day_of_week as usize].push(0);
+				}
 			}
 
 			current_day += Duration::days(1);
 			day_of_week = current_day.weekday().num_days_from_monday() % 7;
 		}
 
-		heatmap
-	}
-
-	fn months_row(&self) -> String {
-		let mut row = "   ".to_string();
-
-		let mut last_index = 0;
-		let mul = match self.format {
-			Format::Chars => 2,
-			Format::Numbers => 3,
-		};
-
-		for (index, month) in &self.months {
-			let range_size = (index * mul)
-				.saturating_sub(last_index * mul)
-				.saturating_sub(3);
-			for _i in 0..range_size {
-				row.push(' ');
-			}
-			last_index = *index;
-			row.push_str(month);
+		if chunk_idx <= months_per_row {
+			heatmap.chunks.push(chunk);
 		}
 
-		row
+		heatmap
 	}
 }
 
@@ -136,29 +126,8 @@ impl Display for Heatmap {
 		writeln!(f, "{} author(s)", authors).unwrap();
 		writeln!(f, "{} commit(s)\n", commits).unwrap();
 
-		writeln!(f, "{}", self.months_row()).unwrap();
-
-		for (day, row) in DAYS.iter().zip(&self.data) {
-			write!(f, "{day} ").unwrap();
-			for val in row {
-				match val {
-					x if *x >= 0 => {
-						let color = &get_color_map()[get_color(*val, self.highest_count)];
-						match self.format {
-							Format::Chars => write!(f, "{color}{}{RESET} ", get_char()).unwrap(),
-							Format::Numbers => {
-								let val = val.min(&99);
-								write!(f, "{color}{:0>2}{RESET} ", val).unwrap();
-							}
-						}
-					}
-					x if *x < 0 => match self.format {
-						Format::Chars => write!(f, "{RESET}  ").unwrap(),
-						Format::Numbers => write!(f, "{RESET}   ").unwrap(),
-					},
-					_ => {}
-				}
-			}
+		for chunk in &self.chunks {
+			chunk.display(self, f);
 			writeln!(f).unwrap();
 		}
 
@@ -172,6 +141,81 @@ impl Display for Heatmap {
 	}
 }
 
+struct Chunk {
+	data: [Vec<i32>; 7],
+	months: Vec<(usize, String)>,
+}
+
+impl Chunk {
+	pub fn new(day_of_week: u32) -> Self {
+		let mut chunk = Self {
+			data: [vec![], vec![], vec![], vec![], vec![], vec![], vec![]],
+			months: vec![],
+		};
+
+		if day_of_week != 0 {
+			for i in 0..day_of_week {
+				chunk.data[i as usize].push(-1);
+			}
+		}
+
+		chunk
+	}
+
+	pub fn display<T: Write>(&self, heatmap: &Heatmap, f: &mut T) {
+		writeln!(f, "{}", self.months_row(heatmap)).unwrap();
+
+		for (day, row) in DAYS.iter().zip(&self.data) {
+			write!(f, "{day} ").unwrap();
+			for val in row {
+				match val {
+					x if *x >= 0 => {
+						let color = &get_color_map()[get_color(*val, heatmap.highest_count)];
+						match heatmap.format {
+							Format::Chars => write!(f, "{color}{}{RESET} ", get_char()).unwrap(),
+							Format::Numbers => {
+								let val = val.min(&99);
+								write!(f, "{color}{:0>2}{RESET} ", val).unwrap();
+							}
+						}
+					}
+					x if *x < 0 => match heatmap.format {
+						Format::Chars => write!(f, "{RESET}  ").unwrap(),
+						Format::Numbers => write!(f, "{RESET}   ").unwrap(),
+					},
+					_ => {}
+				}
+			}
+			writeln!(f).unwrap();
+		}
+	}
+
+	fn months_row(&self, heatmap: &Heatmap) -> String {
+		let mut row = "   ".to_string();
+
+		let mut last_index = 0;
+		let mul = match heatmap.format {
+			Format::Chars => 2,
+			Format::Numbers => 3,
+		};
+
+		for (index, month) in &self.months {
+			let range_size = (index * mul)
+				.saturating_sub(last_index * mul)
+				.saturating_sub(3);
+
+			for _i in 0..range_size {
+				row.push(' ');
+			}
+
+			last_index = *index;
+			row.push_str(month);
+		}
+
+		row
+	}
+}
+
 #[derive(Clone, Debug, PartialEq, Eq, ValueEnum)]
 pub enum HeatmapColors {
 	Green,

src/main.rs

@@ -85,11 +85,20 @@ fn main() -> Result<()> {
 	let until = NaiveDate::parse_from_str(&until, "%Y-%m-%d").unwrap();
 
 	let split_months = args.split_months;
+	let months_per_row = args.months_per_row;
 	let format = args.format.clone();
 
 	let commits = get_commits(args, since, until).with_context(|| "Could not fetch commit list")?;
 
-	let heatmap = Heatmap::new(since, until, commits.0, commits.1, split_months, format);
+	let heatmap = Heatmap::new(
+		since,
+		until,
+		commits.0,
+		commits.1,
+		split_months,
+		months_per_row,
+		format,
+	);
 
 	println!("{heatmap}");