Added support for multiple branches and start/end dates

src/cli.rs

@@ -1,5 +1,6 @@
 use std::path::PathBuf;
 
+use chrono::{Duration, Local};
 use clap::{arg, Parser, ValueHint};
 
 use crate::heatmap::HeatmapColors;
@@ -18,6 +19,18 @@ pub struct CliArgs {
 
 	#[arg(long("color"), value_enum, default_value_t = HeatmapColors::Green)]
 	pub color_scheme: HeatmapColors,
-	// #[arg(short, long, default_value = "24")]
-	// pub split: u32,
+
+	#[arg(short, long, num_args(0..))]
+	pub branches: Option<Vec<String>>,
+
+	#[arg(long("since"), default_value_t = get_since_date())]
+	pub since: String,
+
+	#[arg(long("until"))]
+	pub until: Option<String>,
+}
+
+fn get_since_date() -> String {
+	let date = Local::now() - Duration::days(365);
+	date.format("%Y-%m-%d").to_string()
 }

src/heatmap.rs

@@ -1,30 +1,28 @@
-use std::collections::BTreeMap;
+use std::{collections::BTreeMap, fmt::Display};
 
-use chrono::{DateTime, Datelike, Duration, Local, Utc};
+use chrono::{Datelike, Duration, NaiveDate};
 use clap::ValueEnum;
 
 use crate::{get_char, get_color, get_color_map, Commit, DAYS, RESET};
 
 pub struct Heatmap {
-	data: [Vec<u32>; 7],
-	start_date: DateTime<Local>,
-	end_date: DateTime<Local>,
+	data: [Vec<i32>; 7],
+	since: NaiveDate,
+	until: NaiveDate,
 	commits: Vec<Commit>,
 	months: Vec<(usize, String)>,
+	highest_count: i32,
 }
 
 impl Heatmap {
-	pub fn new(
-		start_date: DateTime<Local>,
-		end_date: DateTime<Local>,
-		commits: Vec<Commit>,
-	) -> Self {
+	pub fn new(since: NaiveDate, until: NaiveDate, commits: Vec<Commit>) -> Self {
 		let mut heatmap = Self {
 			data: [vec![], vec![], vec![], vec![], vec![], vec![], vec![]],
-			start_date,
-			end_date,
+			since,
+			until,
 			commits,
 			months: vec![],
+			highest_count: 0,
 		};
 
 		let mut grouped_commits = BTreeMap::new();
@@ -35,19 +33,19 @@ impl Heatmap {
 			*record += 1;
 		}
 
-		let mut current_day = start_date;
-		let mut day_of_week = (current_day.weekday().num_days_from_monday()) % 7;
+		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(0);
+				heatmap.data[i as usize].push(-1);
 			}
 		}
 
-		while current_day <= end_date {
+		while current_day <= until {
 			let month_name = current_day.format("%b").to_string();
 
-			if current_day == start_date {
+			if current_day == since {
 				heatmap.months.push((0, month_name));
 			}
 			else if current_day.day0() == 0 {
@@ -56,12 +54,17 @@ impl Heatmap {
 					.push((heatmap.data[day_of_week as usize].len(), month_name));
 			}
 
-			let value = grouped_commits.get(&current_day.date_naive());
+			let value = grouped_commits.get(&current_day);
 			match value {
-				Some(val) => heatmap.data[day_of_week as usize].push(*val),
+				Some(val) => {
+					heatmap.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),
 			}
-			// println!("{} {value:?}", current_day.date_naive());
+			// println!("{} {value:?}", current_day);
 			current_day += Duration::days(1);
 			day_of_week = current_day.weekday().num_days_from_monday() % 7;
 		}
@@ -85,28 +88,42 @@ impl Heatmap {
 
 		row
 	}
+}
+
+impl Display for Heatmap {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		let start_date = self.since.format("%Y-%b-%d").to_string();
+		let end_date = self.until.format("%Y-%b-%d").to_string();
+		let commits = self.commits.len().to_string();
+
+		write!(f, "{} - {}\n", start_date, end_date).unwrap();
+		write!(f, "{} commits\n\n", commits).unwrap();
+		write!(f, "{}\n", self.months_row()).unwrap();
 
-	pub fn print(&self) {
-		println!(
-			"{} - {}",
-			self.start_date.format("%Y-%b-%d"),
-			self.end_date.format("%Y-%b-%d")
-		);
-		println!("{} commits\n", &self.commits.len());
-		println!("{}", self.months_row());
 		for (day, row) in DAYS.iter().zip(&self.data) {
-			print!("{day} ");
+			write!(f, "{day} ").unwrap();
 			for val in row {
-				let color = &get_color_map()[get_color(*val)];
-				print!("{color}{}{RESET} ", get_char());
+				match val {
+					x if *x >= 0 => {
+						let color = &get_color_map()[get_color(*val, self.highest_count)];
+						write!(f, "{color}{}{RESET} ", get_char()).unwrap();
+					}
+					x if *x < 0 => {
+						write!(f, "{RESET}  ").unwrap();
+					}
+					_ => {}
+				}
 			}
-			print!("\n");
+			write!(f, "\n").unwrap();
 		}
-		print!("\nLess ");
+
+		write!(f, "\nLess ").unwrap();
 		for color in get_color_map() {
-			print!("{color}{}{RESET} ", get_char())
+			write!(f, "{color}{}{RESET} ", get_char()).unwrap();
 		}
-		print!(" More\n");
+		write!(f, " More\n").unwrap();
+
+		Ok(())
 	}
 }
 

src/main.rs

@@ -2,13 +2,14 @@
 #![feature(let_chains)]
 #![allow(dead_code)]
 
-use std::{cmp::Reverse, sync::OnceLock};
+use std::{cmp::Reverse, collections::HashSet, sync::OnceLock};
 
 use anyhow::{Context, Result};
-use chrono::{Date, DateTime, Datelike, Duration, Local, NaiveDate, Offset, TimeZone, Utc};
+use chrono::{DateTime, Duration, Local, NaiveDate, TimeZone};
 use clap::Parser;
 use gix::{bstr::ByteSlice, ObjectId, Repository};
 use heatmap::HeatmapColors;
+use itertools::Itertools;
 use rgb::Rgb;
 
 use crate::{cli::CliArgs, heatmap::Heatmap};
@@ -28,7 +29,7 @@ const GREEN_COLOR_MAP: [Rgb; 5] = [
 	Rgb(14, 68, 41),
 	Rgb(0, 109, 50),
 	Rgb(38, 166, 65),
-	Rgb(57, 211, 83),
+	Rgb(25, 255, 64),
 ];
 
 const RED_COLOR_MAP: [Rgb; 5] = [
@@ -36,9 +37,10 @@ const RED_COLOR_MAP: [Rgb; 5] = [
 	Rgb(208, 169, 35),
 	Rgb(208, 128, 35),
 	Rgb(208, 78, 35),
-	Rgb(208, 35, 64),
+	Rgb(255, 0, 0),
 ];
 
+#[derive(PartialEq, Eq, Hash)]
 struct Commit {
 	id: ObjectId,
 	title: String,
@@ -64,25 +66,39 @@ fn main() -> Result<()> {
 
 	let repo = gix::open(&args.input).unwrap();
 
-	let end_date = Local::now();
-	let start_date = end_date - Duration::days(365);
+	let since = NaiveDate::parse_from_str(&args.since, "%Y-%m-%d").unwrap();
 
-	let commits =
-		get_commits(repo, args, start_date).with_context(|| "Could not fetch commit list")?;
+	let until = args
+		.until
+		.clone()
+		.unwrap_or_else(|| get_default_until(since));
+	let until = NaiveDate::parse_from_str(&until, "%Y-%m-%d").unwrap();
 
-	let heatmap = Heatmap::new(start_date, end_date, commits);
-	heatmap.print();
+	let commits = get_commits(repo, args, since).with_context(|| "Could not fetch commit list")?;
+
+	let heatmap = Heatmap::new(since, until, commits);
+
+	println!("{heatmap}");
 
 	Ok(())
 }
 
-fn get_color(val: u32) -> usize {
-	match val {
-		0 => 0,
-		x if x < 2 => 1,
-		x if x < 4 => 2,
-		x if x < 6 => 3,
-		x if x > 8 => 4,
+fn get_default_until(since: NaiveDate) -> String {
+	let mut until = Local::now().date_naive();
+	if since + Duration::days(365) < until {
+		until = since + Duration::days(365);
+	}
+	until.format("%Y-%m-%d").to_string()
+}
+
+fn get_color(val: i32, high: i32) -> usize {
+	let color = val as f32 / high as f32;
+	match color {
+		0.0 => 0,
+		x if x <= 0.2 => 1,
+		x if x <= 0.4 => 2,
+		x if x <= 0.8 => 3,
+		x if x > 0.8 => 4,
 		_ => 0,
 	}
 }
@@ -106,110 +122,71 @@ fn get_color_map() -> Vec<String> {
 		.to_vec()
 }
 
-fn get_commits(
-	repo: Repository,
-	args: CliArgs,
-	start_date: DateTime<Local>,
-) -> Result<Vec<Commit>> {
-	let mut commits: Vec<Commit> = repo
-		.head()?
-		.into_peeled_id()?
-		.ancestors()
-		.all()?
-		.filter_map(|c| c.ok())
-		.filter_map(|c| c.object().ok())
-		.filter_map(|c| {
-			let title = c
-				.message()
-				.ok()?
-				.title
-				.trim_ascii()
-				.to_str()
-				.ok()?
-				.to_string();
+fn get_commits(repo: Repository, args: CliArgs, start_date: NaiveDate) -> Result<Vec<Commit>> {
+	let mut commits: HashSet<Commit> = HashSet::new();
 
-			let author = c.author().ok()?.name.to_string();
-			if author != args.author {
-				return None;
-			}
+	let branches = match args.branches {
+		Some(b) => b,
+		None => repo
+			.branch_names()
+			.into_iter()
+			.map(|b| b.to_string())
+			.collect_vec(),
+	};
 
-			let time = c.time().ok()?;
-			// let offset = Local.timestamp_opt(0, 0).unwrap().offset().fix();
-			// let time = Local::now() - Duration::seconds(time.seconds);
-			let time = DateTime::from_timestamp_millis(time.seconds * 1000)?.with_timezone(&Local);
-			if time <= start_date + Duration::days(1) {
-				return None;
-			}
+	for branch in branches {
+		let branch_commits = repo
+			.rev_parse(&*branch)?
+			.single()
+			.unwrap()
+			.ancestors()
+			.all()?;
 
-			Some(Commit {
-				id: c.id,
-				title,
-				author,
-				time,
+		branch_commits
+			.filter_map(|c| c.ok())
+			.filter_map(|c| c.object().ok())
+			.filter_map(|c| {
+				let title = c
+					.message()
+					.ok()?
+					.title
+					.trim_ascii()
+					.to_str()
+					.ok()?
+					.to_string();
+
+				let author = c.author().ok()?.name.to_string();
+				if author != args.author {
+					return None;
+				}
+
+				let current_time = Local::now().time();
+				let start_date = start_date.and_time(current_time);
+				let start_date = Local.from_local_datetime(&start_date).unwrap();
+
+				let time = c.time().ok()?;
+				let time =
+					DateTime::from_timestamp_millis(time.seconds * 1000)?.with_timezone(&Local);
+				if time <= start_date + Duration::days(1) {
+					return None;
+				}
+
+				Some(Commit {
+					id: c.id,
+					title,
+					author,
+					time,
+				})
 			})
-		})
-		.collect();
+			.for_each(|c| {
+				commits.insert(c);
+			});
+	}
 
-	commits.sort_by_cached_key(|a| Reverse(a.time));
+	let commits = commits
+		.into_iter()
+		.sorted_by_cached_key(|a| Reverse(a.time))
+		.collect_vec();
 
 	Ok(commits)
 }
-
-// fn print_streak(commits: Vec<Commit>) {
-// 	let mut commits_pushed = 0;
-// 	let mut days_streak = 0;
-// 	let mut last_date: Option<DateTime<Utc>> = None;
-// 	let mut start_date: Option<DateTime<Utc>> = None;
-// 	let mut end_date: Option<DateTime<Utc>> = None;
-//
-// 	for commit in commits {
-// 		match last_date {
-// 			Some(date) => {
-// 				let day = date.ordinal0();
-// 				let commit_day = commit.time.ordinal0();
-// 				let time_diff = date - commit.time;
-//
-// 				if commit_day != day {
-// 					days_streak += 1;
-// 				}
-//
-// 				if time_diff.num_hours() >= 24 {
-// 					// println!(
-// 					// 	"Failing Commit\n{} - {} {} hours difference",
-// 					// 	commit.id,
-// 					// 	commit.time.to_rfc3339(),
-// 					// 	time_diff.num_hours()
-// 					// );
-// 					break;
-// 				}
-//
-// 				commits_pushed += 1;
-// 				last_date = Some(commit.time);
-// 				end_date = Some(commit.time);
-//
-// 				// println!(
-// 				// 	"{} - {} {} hours difference",
-// 				// 	commit.id,
-// 				// 	commit.time.to_rfc3339(),
-// 				// 	time_diff.num_hours()
-// 				// );
-// 			}
-// 			None => {
-// 				last_date = Some(commit.time);
-// 				start_date = Some(commit.time);
-// 				// println!("First Commit\n{} - {}", commit.id, commit.time.to_rfc3339());
-// 				continue;
-// 			}
-// 		}
-// 	}
-//
-// 	println!("{commits_pushed} commits pushed during a {days_streak} days streak");
-// 	if let Some(start_date) = start_date
-// 		&& let Some(end_date) = end_date
-// 	{
-// 		let start_date = start_date.format("%d-%b-%Y").to_string();
-// 		let end_date = end_date.format("%d-%b-%Y").to_string();
-//
-// 		println!("Between: {start_date} {end_date}");
-// 	}
-// }