use chrono::NaiveDate;
use chrono::Weekday;
use std::collections::HashSet;
#[derive(PartialEq, Eq, Clone, Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct Calendar {
weekend: HashSet<Weekday>,
holidays: HashSet<NaiveDate>,
}
impl Default for Calendar {
fn default() -> Self {
Self::new()
}
}
pub fn basic_calendar() -> Calendar {
Calendar {
weekend: [Weekday::Sat, Weekday::Sun].into_iter().collect(),
holidays: HashSet::new(),
}
}
pub fn calendar_unions(calendars: &[Calendar]) -> Calendar {
let mut result = Calendar::new();
for cal in calendars {
result.union(cal);
}
result
}
impl Calendar {
pub fn new() -> Self {
Self {
weekend: HashSet::new(),
holidays: HashSet::new(),
}
}
pub fn get_holidays(&self) -> &HashSet<NaiveDate> {
&self.holidays
}
pub fn get_weekend(&self) -> &HashSet<Weekday> {
&self.weekend
}
pub fn add_holidays(&mut self, holidays: &HashSet<NaiveDate>) {
self.holidays = self.holidays.union(holidays).cloned().collect();
}
pub fn add_weekends(&mut self, weekends: &HashSet<Weekday>) {
self.weekend = self.weekend.union(weekends).cloned().collect();
}
pub fn union(&mut self, other: &Calendar) {
self.holidays = self.holidays.union(&other.holidays).cloned().collect();
self.weekend = self.weekend.union(&other.weekend).cloned().collect();
}
pub fn is_business_day(&self, date: &chrono::NaiveDate) -> bool {
crate::algebra::is_business_day(date, self)
}
pub fn intersection(&mut self, other: &Calendar) {
self.holidays = self
.holidays
.intersection(&other.holidays)
.cloned()
.collect();
self.weekend = self.weekend.intersection(&other.weekend).cloned().collect();
}
}
#[cfg(test)]
mod tests {
use crate::calendar::{self as c, Calendar};
use chrono::{NaiveDate, Weekday};
use std::collections::HashSet;
#[test]
fn add_holidays_test() {
let mut cal = c::basic_calendar();
let christmas_day = NaiveDate::from_ymd_opt(2023, 12, 25).unwrap();
let boxing_day = NaiveDate::from_ymd_opt(2023, 12, 26).unwrap();
let new_holidays: HashSet<NaiveDate> = [christmas_day, boxing_day].into_iter().collect();
cal.add_holidays(&new_holidays);
assert_eq!(cal.holidays, new_holidays);
}
#[test]
fn add_weekends_test() {
let mut cal = Calendar::new();
let new_weekend: HashSet<Weekday> = [Weekday::Mon].into_iter().collect();
cal.add_weekends(&new_weekend);
assert_eq!(cal.weekend, new_weekend);
}
#[test]
fn get_holidays_test() {
let mut cal = c::basic_calendar();
let christmas_day = NaiveDate::from_ymd_opt(2023, 12, 25).unwrap();
let boxing_day = NaiveDate::from_ymd_opt(2023, 12, 26).unwrap();
let new_holidays: HashSet<NaiveDate> = [christmas_day, boxing_day].into_iter().collect();
cal.add_holidays(&new_holidays);
assert_eq!(cal.get_holidays(), &new_holidays);
}
#[test]
fn get_weekend_test() {
let mut cal = Calendar::new();
let new_weekend: HashSet<Weekday> = [Weekday::Mon].into_iter().collect();
cal.add_weekends(&new_weekend);
assert_eq!(cal.get_weekend(), &new_weekend);
}
#[test]
fn calendar_union_test() {
let christmas_day = NaiveDate::from_ymd_opt(2023, 12, 25).unwrap();
let boxing_day = NaiveDate::from_ymd_opt(2023, 12, 26).unwrap();
let mut cal1 = Calendar::new();
cal1.add_weekends(&[Weekday::Sat].into_iter().collect());
cal1.add_holidays(&[christmas_day].into_iter().collect());
let mut cal2 = Calendar::new();
cal2.add_weekends(&[Weekday::Sun].into_iter().collect());
cal2.add_holidays(&[boxing_day].into_iter().collect());
let mut expected = c::basic_calendar();
expected.add_holidays(&[christmas_day, boxing_day].into_iter().collect());
cal1.union(&cal2);
assert_eq!(cal1, expected);
}
#[test]
fn calendar_intersection_test() {
let christmas_day = NaiveDate::from_ymd_opt(2023, 12, 25).unwrap();
let boxing_day = NaiveDate::from_ymd_opt(2023, 12, 26).unwrap();
let mut cal1 = Calendar::new();
cal1.add_weekends(&[Weekday::Sun].into_iter().collect());
cal1.add_holidays(&[christmas_day].into_iter().collect());
let mut cal2 = Calendar::new();
cal2.add_weekends(&[Weekday::Sun].into_iter().collect());
cal2.add_holidays(&[christmas_day, boxing_day].into_iter().collect());
let mut expected = Calendar::new();
expected.add_weekends(&[Weekday::Sun].into_iter().collect());
expected.add_holidays(&[christmas_day].into_iter().collect());
cal1.intersection(&cal2);
assert_eq!(cal1, expected);
}
#[test]
fn default_is_empty_test() {
let cal = Calendar::default();
assert!(cal.get_holidays().is_empty());
assert!(cal.get_weekend().is_empty());
}
#[test]
fn calendar_unions_test() {
let xmas = NaiveDate::from_ymd_opt(2024, 12, 25).unwrap();
let mut cal1 = Calendar::new();
cal1.add_weekends(&[Weekday::Sat].into_iter().collect());
let mut cal2 = Calendar::new();
cal2.add_holidays(&[xmas].into_iter().collect());
let combined = c::calendar_unions(&[cal1, cal2]);
assert!(combined.get_weekend().contains(&Weekday::Sat));
assert!(combined.get_holidays().contains(&xmas));
}
}