magnus167/rustframe
1
0
Fork 0
mirror of https://github.com/Magnus167/rustframe.git synced 2025-08-20 04:00:01 +00:00
Code Issues Packages Projects Releases Wiki Activity
rustframe/src/utils/dateutils/dates.rs
Palash Tyagi 52729e55be Clean up comments and improve code clarity
2025-05-15 00:00:56 +01:00

2172 lines
79 KiB
Rust
Raw Blame History

use chrono::{Datelike, Duration, NaiveDate, Weekday};
use std::collections::HashMap;
use std::error::Error;
use std::hash::Hash;
use std::result::Result;
use std::str::FromStr;
/// Represents the frequency at which calendar dates should be generated.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum DateFreq {
Daily, // Every calendar day
WeeklyMonday, // Every Monday
WeeklyFriday, // Every Friday
MonthStart, // First calendar day of the month
MonthEnd, // Last calendar day of the month
QuarterStart, // First calendar day of the quarter
QuarterEnd, // Last calendar day of the quarter
YearStart, // First calendar day of the year (Jan 1st)
YearEnd, // Last calendar day of the year (Dec 31st)
}
/// Indicates whether the first or last date in a periodic group (like month, quarter)
/// is selected for the frequency.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AggregationType {
Start, // Indicates picking the first calendar date in a group's period.
End, // Indicates picking the last calendar day in a group's period.
}
impl DateFreq {
/// Attempts to parse a frequency string into a `DateFreq` enum.
///
/// This is a convenience wrapper around `from_str`.
///
/// # Arguments
///
/// * `freq` - The frequency string (e.g., "D", "W", "ME").
///
/// # Errors
///
/// Returns an error if the string does not match any known frequency.
pub fn from_string(freq: String) -> Result<Self, Box<dyn Error>> {
freq.parse()
}
/// Returns the canonical string representation of the frequency.
///
/// This returns the primary code (e.g., "D", "W", "Y", "YE"), not the aliases.
pub fn to_string(&self) -> String {
let r = match self {
DateFreq::Daily => "D",
DateFreq::WeeklyMonday => "W",
DateFreq::MonthStart => "M",
DateFreq::QuarterStart => "Q",
DateFreq::YearStart => "Y",
DateFreq::MonthEnd => "ME",
DateFreq::QuarterEnd => "QE",
DateFreq::WeeklyFriday => "WF",
DateFreq::YearEnd => "YE",
};
r.to_string()
}
/// Determines whether the frequency represents a start-of-period or end-of-period aggregation.
pub fn agg_type(&self) -> AggregationType {
match self {
DateFreq::Daily
| DateFreq::WeeklyMonday
| DateFreq::MonthStart
| DateFreq::QuarterStart
| DateFreq::YearStart => AggregationType::Start,
DateFreq::WeeklyFriday
| DateFreq::MonthEnd
| DateFreq::QuarterEnd
| DateFreq::YearEnd => AggregationType::End,
}
}
}
// Implement FromStr for DateFreq to allow parsing directly using `parse()`
impl FromStr for DateFreq {
type Err = Box<dyn Error>;
/// Attempts to parse a frequency string slice into a `DateFreq` enum.
///
/// Supports various frequency codes and common aliases.
///
/// | Code | Alias | Description |
/// |------|---------|--------------------------|
/// | D | | Daily (every day) |
/// | W | WS | Weekly Monday |
/// | M | MS | Month Start (1st) |
/// | Q | QS | Quarter Start (1st) |
/// | Y | A, AS, YS | Year Start (Jan 1st) |
/// | ME | | Month End (Last day) |
/// | QE | | Quarter End (Last day) |
/// | WF | | Weekly Friday |
/// | YE | AE | Year End (Dec 31st) |
///
/// # Arguments
///
/// * `freq` - The frequency string slice (e.g., "D", "W", "ME").
///
/// # Errors
///
/// Returns an error if the string does not match any known frequency.
fn from_str(freq: &str) -> Result<Self, Self::Err> {
let r = match freq {
"D" => DateFreq::Daily,
"W" | "WS" => DateFreq::WeeklyMonday,
"M" | "MS" => DateFreq::MonthStart,
"Q" | "QS" => DateFreq::QuarterStart,
"Y" | "A" | "AS" | "YS" => DateFreq::YearStart,
"ME" => DateFreq::MonthEnd,
"QE" => DateFreq::QuarterEnd,
"WF" => DateFreq::WeeklyFriday,
"YE" | "AE" => DateFreq::YearEnd,
_ => return Err(format!("Invalid frequency specified: {}", freq).into()),
};
Ok(r)
}
}
/// Represents a list of calendar dates generated between a start and end date
/// at a specified frequency. Provides methods to retrieve the full list,
/// count, or dates grouped by period.
#[derive(Debug, Clone)]
pub struct DatesList {
start_date_str: String,
end_date_str: String,
freq: DateFreq,
}
// Helper enum to represent the key for grouping dates into periods.
// Deriving traits for comparison and hashing allows using it as a HashMap key
// and for sorting groups chronologically.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
enum GroupKey {
Daily(NaiveDate), // Group by the specific date (for Daily frequency)
Weekly(i32, u32), // Group by year and ISO week number
Monthly(i32, u32), // Group by year and month (1-12)
Quarterly(i32, u32), // Group by year and quarter (1-4)
Yearly(i32), // Group by year
}
/// Represents a collection of calendar dates generated according to specific rules.
///
/// It can be defined either by a start and end date range or by a start date
/// and a fixed number of periods. It provides methods to retrieve the dates
/// as a flat list, count them, or group them by their natural period
/// (e.g., month, quarter).
///
/// This struct handles all calendar dates, including weekends.
///
/// ## Examples
///
/// **1. Using `DatesList::new` (Start and End Date):**
///
/// ```rust
/// use chrono::NaiveDate;
/// use std::error::Error;
/// # use rustframe::utils::{DatesList, DateFreq}; // Assuming the crate/module is named 'dates'
///
/// # fn main() -> Result<(), Box<dyn Error>> {
/// let start_date = "2023-11-01".to_string(); // Wednesday
/// let end_date = "2023-11-07".to_string(); // Tuesday
/// let freq = DateFreq::Daily;
///
/// let dates_list = DatesList::new(start_date, end_date, freq);
///
/// let expected_dates = vec![
/// NaiveDate::from_ymd_opt(2023, 11, 1).unwrap(), // Wed
/// NaiveDate::from_ymd_opt(2023, 11, 2).unwrap(), // Thu
/// NaiveDate::from_ymd_opt(2023, 11, 3).unwrap(), // Fri
/// NaiveDate::from_ymd_opt(2023, 11, 4).unwrap(), // Sat
/// NaiveDate::from_ymd_opt(2023, 11, 5).unwrap(), // Sun
/// NaiveDate::from_ymd_opt(2023, 11, 6).unwrap(), // Mon
/// NaiveDate::from_ymd_opt(2023, 11, 7).unwrap(), // Tue
/// ];
///
/// assert_eq!(dates_list.list()?, expected_dates);
/// assert_eq!(dates_list.count()?, 7);
/// # Ok(())
/// # }
/// ```
///
/// **2. Using `DatesList::from_n_periods` (Start Date and Count):**
///
/// ```rust
/// use chrono::NaiveDate;
/// use std::error::Error;
/// use rustframe::utils::{DatesList, DateFreq};
///
/// fn main() -> Result<(), Box<dyn Error>> {
/// let start_date = "2024-02-28".to_string(); // Wednesday
/// let freq = DateFreq::WeeklyFriday;
/// let n_periods = 3;
///
/// let dates_list = DatesList::from_n_periods(start_date, freq, n_periods)?;
///
/// // The first Friday on or after 2024-02-28 is Mar 1.
/// // The next two Fridays are Mar 8 and Mar 15.
/// let expected_dates = vec![
/// NaiveDate::from_ymd_opt(2024, 3, 1).unwrap(),
/// NaiveDate::from_ymd_opt(2024, 3, 8).unwrap(),
/// NaiveDate::from_ymd_opt(2024, 3, 15).unwrap(),
/// ];
///
/// assert_eq!(dates_list.list()?, expected_dates);
/// assert_eq!(dates_list.count()?, 3);
/// assert_eq!(dates_list.start_date_str(), "2024-02-28"); // Keeps original start string
/// assert_eq!(dates_list.end_date_str(), "2024-03-15"); // End date is the last generated date
/// # Ok(())
/// # }
/// ```
///
/// **3. Using `groups()`:**
///
/// ```rust
/// use chrono::NaiveDate;
/// use std::error::Error;
/// use rustframe::utils::{DatesList, DateFreq};
///
/// fn main() -> Result<(), Box<dyn Error>> {
/// let start_date = "2023-11-20".to_string(); // Mon, Week 47
/// let end_date = "2023-12-08".to_string(); // Fri, Week 49
/// let freq = DateFreq::MonthEnd; // Find month-ends
///
/// let dates_list = DatesList::new(start_date, end_date, freq);
///
/// // Month ends >= Nov 20 and <= Dec 08: Nov 30
/// let groups = dates_list.groups()?;
///
/// assert_eq!(groups.len(), 1); // Only November's end date falls in the range
/// assert_eq!(groups[0], vec![NaiveDate::from_ymd_opt(2023, 11, 30).unwrap()]); // Nov 2023 group
/// Ok(())
/// }
/// ```
impl DatesList {
/// Creates a new `DatesList` instance defined by a start and end date.
///
/// # Arguments
///
/// * `start_date_str` - The inclusive start date as a string (e.g., "YYYY-MM-DD").
/// * `end_date_str` - The inclusive end date as a string (e.g., "YYYY-MM-DD").
/// * `freq` - The frequency for generating dates.
pub fn new(start_date_str: String, end_date_str: String, freq: DateFreq) -> Self {
DatesList {
start_date_str,
end_date_str,
freq,
}
}
/// Creates a new `DatesList` instance defined by a start date, frequency,
/// and the number of periods (dates) to generate.
///
/// This calculates the required dates using a `DatesGenerator` and determines
/// the effective end date based on the last generated date.
///
/// # Arguments
///
/// * `start_date_str` - The start date as a string (e.g., "YYYY-MM-DD"). The first generated date will be on or after this date.
/// * `freq` - The frequency for generating dates.
/// * `n_periods` - The exact number of dates to generate according to the frequency.
///
/// # Errors
///
/// Returns an error if:
/// * `start_date_str` cannot be parsed.
/// * `n_periods` is 0 (as this would result in an empty list and no defined end date).
pub fn from_n_periods(
start_date_str: String,
freq: DateFreq,
n_periods: usize,
) -> Result<Self, Box<dyn Error>> {
if n_periods == 0 {
return Err("n_periods must be greater than 0".into());
}
let start_date = NaiveDate::parse_from_str(&start_date_str, "%Y-%m-%d")?;
// Use the generator to find all the dates
let generator = DatesGenerator::new(start_date, freq, n_periods)?;
let dates: Vec<NaiveDate> = generator.collect();
// Should always have at least one date if n_periods > 0 and generator construction succeeded
let last_date = dates
.last()
.ok_or("Generator failed to produce dates even though n_periods > 0")?;
let end_date_str = last_date.format("%Y-%m-%d").to_string();
Ok(DatesList {
start_date_str, // Keep the original start date string
end_date_str,
freq,
})
}
/// Returns the flat list of dates within the specified range and frequency.
///
/// The list is guaranteed to be sorted chronologically.
///
/// # Errors
///
/// Returns an error if the start or end date strings cannot be parsed.
pub fn list(&self) -> Result<Vec<NaiveDate>, Box<dyn Error>> {
// Delegate the core logic to the internal helper function
get_dates_list_with_freq(&self.start_date_str, &self.end_date_str, self.freq)
}
/// Returns the count of dates within the specified range and frequency.
///
/// # Errors
///
/// Returns an error if the start or end date strings cannot be parsed (as it
/// calls `list` internally).
pub fn count(&self) -> Result<usize, Box<dyn Error>> {
self.list().map(|list| list.len())
}
/// Returns a list of date lists, where each inner list contains dates
/// belonging to the same period (determined by frequency).
///
/// The outer list (groups) is sorted by period chronologically, and the
/// inner lists (dates within groups) are also sorted chronologically.
///
/// For `Daily` frequency, each date forms its own group. For `Weekly`
/// frequencies, grouping is by ISO week number. For `Monthly`, `Quarterly`,
/// and `Yearly` frequencies, grouping is by the respective period.
///
/// # Errors
///
/// Returns an error if the start or end date strings cannot be parsed.
pub fn groups(&self) -> Result<Vec<Vec<NaiveDate>>, Box<dyn Error>> {
let dates = self.list()?;
group_dates_helper(dates, self.freq)
}
/// Returns the start date parsed as a `NaiveDate`.
///
/// # Errors
///
/// Returns an error if the start date string is not in "YYYY-MM-DD" format.
pub fn start_date(&self) -> Result<NaiveDate, Box<dyn Error>> {
NaiveDate::parse_from_str(&self.start_date_str, "%Y-%m-%d").map_err(|e| e.into())
}
/// Returns the start date string.
pub fn start_date_str(&self) -> &str {
&self.start_date_str
}
/// Returns the end date parsed as a `NaiveDate`.
///
/// # Errors
///
/// Returns an error if the end date string is not in "YYYY-MM-DD" format.
pub fn end_date(&self) -> Result<NaiveDate, Box<dyn Error>> {
NaiveDate::parse_from_str(&self.end_date_str, "%Y-%m-%d").map_err(|e| e.into())
}
/// Returns the end date string.
pub fn end_date_str(&self) -> &str {
&self.end_date_str
}
/// Returns the frequency enum.
pub fn freq(&self) -> DateFreq {
self.freq
}
/// Returns the canonical string representation of the frequency.
pub fn freq_str(&self) -> String {
self.freq.to_string()
}
}
/// An iterator that generates a sequence of calendar dates based on a start date,
/// frequency, and a specified number of periods.
///
/// This implements the `Iterator` trait, allowing generation of dates one by one.
/// It's useful when you need to process dates lazily or only need a fixed number
/// starting from a specific point, without necessarily defining an end date beforehand.
/// # Examples
///
/// **1. Basic Iteration (Month End):**
///
/// ```rust
/// use chrono::NaiveDate;
/// use std::error::Error;
/// use rustframe::utils::{DatesGenerator, DateFreq};
///
/// fn main() -> Result<(), Box<dyn Error>> {
/// let start = NaiveDate::from_ymd_opt(2023, 12, 28).unwrap(); // Thursday
/// let freq = DateFreq::MonthEnd;
/// let n_periods = 4; // Dec '23, Jan '24, Feb '24, Mar '24
///
/// let mut generator = DatesGenerator::new(start, freq, n_periods)?;
///
/// // First month-end on or after 2023-12-28 is 2023-12-31
/// assert_eq!(generator.next(), Some(NaiveDate::from_ymd_opt(2023, 12, 31).unwrap()));
/// assert_eq!(generator.next(), Some(NaiveDate::from_ymd_opt(2024, 1, 31).unwrap()));
/// assert_eq!(generator.next(), Some(NaiveDate::from_ymd_opt(2024, 2, 29).unwrap())); // Leap year
/// assert_eq!(generator.next(), Some(NaiveDate::from_ymd_opt(2024, 3, 31).unwrap()));
/// assert_eq!(generator.next(), None); // Exhausted
/// Ok(())
/// }
/// ```
///
/// **2. Collecting into a Vec (Daily):**
///
/// ```rust
/// use chrono::NaiveDate;
/// use std::error::Error;
/// use rustframe::utils::{DatesGenerator, DateFreq};
///
/// fn main() -> Result<(), Box<dyn Error>> {
/// let start = NaiveDate::from_ymd_opt(2024, 4, 29).unwrap(); // Monday
/// let freq = DateFreq::Daily;
/// let n_periods = 5;
///
/// let generator = DatesGenerator::new(start, freq, n_periods)?;
/// let dates: Vec<NaiveDate> = generator.collect();
///
/// let expected_dates = vec![
/// NaiveDate::from_ymd_opt(2024, 4, 29).unwrap(), // Mon
/// NaiveDate::from_ymd_opt(2024, 4, 30).unwrap(), // Tue
/// NaiveDate::from_ymd_opt(2024, 5, 1).unwrap(), // Wed
/// NaiveDate::from_ymd_opt(2024, 5, 2).unwrap(), // Thu
/// NaiveDate::from_ymd_opt(2024, 5, 3).unwrap(), // Fri
/// ];
///
/// assert_eq!(dates, expected_dates);
/// Ok(())
/// }
/// ```
///
/// **3. Starting on the Exact Day (Weekly Monday):**
///
/// ```rust
/// use chrono::NaiveDate;
/// use std::error::Error;
/// use rustframe::utils::{DatesGenerator, DateFreq};
///
/// fn main() -> Result<(), Box<dyn Error>> {
/// let start = NaiveDate::from_ymd_opt(2024, 1, 1).unwrap(); // Monday
/// let freq = DateFreq::WeeklyMonday;
/// let n_periods = 3;
///
/// let mut generator = DatesGenerator::new(start, freq, n_periods)?;
///
/// assert_eq!(generator.next(), Some(NaiveDate::from_ymd_opt(2024, 1, 1).unwrap()));
/// assert_eq!(generator.next(), Some(NaiveDate::from_ymd_opt(2024, 1, 8).unwrap()));
/// assert_eq!(generator.next(), Some(NaiveDate::from_ymd_opt(2024, 1, 15).unwrap()));
/// assert_eq!(generator.next(), None);
/// Ok(())
/// }
/// ```
#[derive(Debug, Clone)]
pub struct DatesGenerator {
pub freq: DateFreq,
pub periods_remaining: usize,
// Stores the *next* date to be yielded by the iterator.
pub next_date_candidate: Option<NaiveDate>,
}
impl DatesGenerator {
/// Creates a new `DatesGenerator`.
///
/// It calculates the first valid date based on the `start_date` and `freq`,
/// which will be the first item yielded by the iterator.
///
/// # Arguments
///
/// * `start_date` - The date from which to start searching for the first valid date.
/// * `freq` - The frequency for generating dates.
/// * `n_periods` - The total number of dates to generate.
///
/// # Errors
///
/// Returns an error if initial date calculation fails (e.g., due to overflow, though unlikely).
pub fn new(
start_date: NaiveDate,
freq: DateFreq,
n_periods: usize,
) -> Result<Self, Box<dyn Error>> {
let first_date = if n_periods > 0 {
Some(find_first_date_on_or_after(start_date, freq)?)
} else {
None // No dates to generate if n_periods is 0
};
Ok(DatesGenerator {
freq,
periods_remaining: n_periods,
next_date_candidate: first_date,
})
}
}
impl Iterator for DatesGenerator {
type Item = NaiveDate;
/// Returns the next date in the sequence, or `None` if `n_periods`
/// dates have already been generated.
fn next(&mut self) -> Option<Self::Item> {
match self.next_date_candidate {
Some(current_date) if self.periods_remaining > 0 => {
// Prepare the *next* candidate for the subsequent call
// We calculate the next date *before* decrementing periods_remaining
// If find_next_date fails, we treat it as the end of the sequence.
self.next_date_candidate = find_next_date(current_date, self.freq).ok();
// Decrement the count *after* potentially getting the next date
self.periods_remaining -= 1;
// Return the stored current date
Some(current_date)
}
_ => {
// Exhausted or no initial date
self.periods_remaining = 0; // Ensure it's 0
self.next_date_candidate = None;
None
}
}
}
}
// Internal helper functions
pub fn group_dates_helper(
dates: Vec<NaiveDate>,
freq: DateFreq,
) -> Result<Vec<Vec<NaiveDate>>, Box<dyn Error + 'static>> {
let mut groups: HashMap<GroupKey, Vec<NaiveDate>> = HashMap::new();
for date in dates {
let key = match freq {
DateFreq::Daily => GroupKey::Daily(date),
DateFreq::WeeklyMonday | DateFreq::WeeklyFriday => {
let iso_week = date.iso_week();
GroupKey::Weekly(iso_week.year(), iso_week.week())
}
DateFreq::MonthStart | DateFreq::MonthEnd => {
GroupKey::Monthly(date.year(), date.month())
}
DateFreq::QuarterStart | DateFreq::QuarterEnd => {
GroupKey::Quarterly(date.year(), month_to_quarter(date.month()))
}
DateFreq::YearStart | DateFreq::YearEnd => GroupKey::Yearly(date.year()),
};
groups.entry(key).or_insert_with(Vec::new).push(date);
}
let mut sorted_groups: Vec<(GroupKey, Vec<NaiveDate>)> = groups.into_iter().collect();
sorted_groups.sort_by(|(k1, _), (k2, _)| k1.cmp(k2));
// Dates within groups are already sorted because they came from the sorted `self.list()`.
let result_groups = sorted_groups.into_iter().map(|(_, dates)| dates).collect();
Ok(result_groups)
}
/// Generates the flat list of dates for the given range and frequency.
/// Assumes the `collect_*` functions return sorted dates.
pub fn get_dates_list_with_freq(
start_date_str: &str,
end_date_str: &str,
freq: DateFreq,
) -> Result<Vec<NaiveDate>, Box<dyn Error>> {
let start_date = NaiveDate::parse_from_str(start_date_str, "%Y-%m-%d")?;
let end_date = NaiveDate::parse_from_str(end_date_str, "%Y-%m-%d")?;
if start_date > end_date {
return Ok(Vec::new());
}
let dates = match freq {
DateFreq::Daily => collect_daily(start_date, end_date)?,
DateFreq::WeeklyMonday => collect_weekly(start_date, end_date, Weekday::Mon)?,
DateFreq::WeeklyFriday => collect_weekly(start_date, end_date, Weekday::Fri)?,
DateFreq::MonthStart => {
collect_monthly(start_date, end_date, /*want_first_day=*/ true)?
}
DateFreq::MonthEnd => {
collect_monthly(start_date, end_date, /*want_first_day=*/ false)?
}
DateFreq::QuarterStart => {
collect_quarterly(start_date, end_date, /*want_first_day=*/ true)?
}
DateFreq::QuarterEnd => {
collect_quarterly(start_date, end_date, /*want_first_day=*/ false)?
}
DateFreq::YearStart => collect_yearly(start_date, end_date, /*want_first_day=*/ true)?,
DateFreq::YearEnd => collect_yearly(start_date, end_date, /*want_first_day=*/ false)?,
};
// The collect_* functions should now generate sorted dates directly.
Ok(dates)
}
// Low-Level Date Collection Functions (Internal)
// These functions generate dates within a *range* [start_date, end_date]
/// Returns all calendar days day-by-day within the range.
fn collect_daily(
start_date: NaiveDate,
end_date: NaiveDate,
) -> Result<Vec<NaiveDate>, Box<dyn Error>> {
let mut result = Vec::new();
let mut current = start_date;
while current <= end_date {
result.push(current);
current = current
.succ_opt()
.ok_or("Date overflow near end of supported range")?;
}
Ok(result)
}
/// Returns the specified `target_weekday` in each week within the range.
fn collect_weekly(
start_date: NaiveDate,
end_date: NaiveDate,
target_weekday: Weekday,
) -> Result<Vec<NaiveDate>, Box<dyn Error>> {
let mut result = Vec::new();
let mut current = move_to_day_of_week_on_or_after(start_date, target_weekday)?;
while current <= end_date {
result.push(current);
current = current
.checked_add_signed(Duration::days(7))
.ok_or("Date overflow adding 7 days")?;
}
Ok(result)
}
/// Returns either the first or last calendar day in each month of the range.
fn collect_monthly(
start_date: NaiveDate,
end_date: NaiveDate,
want_first_day: bool,
) -> Result<Vec<NaiveDate>, Box<dyn Error>> {
let mut result = Vec::new();
let mut year = start_date.year();
let mut month = start_date.month();
let next_month = |(yr, mo): (i32, u32)| -> (i32, u32) {
if mo == 12 {
(yr + 1, 1)
} else {
(yr, mo + 1)
}
};
loop {
let candidate = if want_first_day {
first_day_of_month(year, month)?
} else {
last_day_of_month(year, month)?
};
if candidate > end_date {
break;
}
if candidate >= start_date {
result.push(candidate);
}
if year > end_date.year() || (year == end_date.year() && month >= end_date.month()) {
break;
}
let (ny, nm) = next_month((year, month));
year = ny;
month = nm;
// Safety check for potential infinite loop, though unlikely with valid date logic
if year > end_date.year() + 2 {
return Err("Loop seems to exceed reasonable year range in collect_monthly".into());
}
}
Ok(result)
}
/// Return either the first or last calendar day in each quarter of the range.
fn collect_quarterly(
start_date: NaiveDate,
end_date: NaiveDate,
want_first_day: bool,
) -> Result<Vec<NaiveDate>, Box<dyn Error>> {
let mut result = Vec::new();
let mut year = start_date.year();
let mut q = month_to_quarter(start_date.month());
loop {
let candidate = if want_first_day {
first_day_of_quarter(year, q)?
} else {
last_day_of_quarter(year, q)?
};
if candidate > end_date {
break;
}
if candidate >= start_date {
result.push(candidate);
}
let end_q = month_to_quarter(end_date.month());
if year > end_date.year() || (year == end_date.year() && q >= end_q) {
break;
}
if q == 4 {
year += 1;
q = 1;
} else {
q += 1;
}
// Safety check
if year > end_date.year() + 2 {
return Err("Loop seems to exceed reasonable year range in collect_quarterly".into());
}
}
Ok(result)
}
/// Returns a list of dates between the given start and end dates, inclusive,
/// at the specified frequency.
/// This function is a convenience wrapper around `get_dates_list_with_freq`.
pub fn get_dates_list_with_freq_from_naive_date(
start_date: NaiveDate,
end_date: NaiveDate,
freq: DateFreq,
) -> Result<Vec<NaiveDate>, Box<dyn Error>> {
get_dates_list_with_freq(
&start_date.format("%Y-%m-%d").to_string(),
&end_date.format("%Y-%m-%d").to_string(),
freq,
)
}
/// Return either the first or last calendar day in each year of the range.
fn collect_yearly(
start_date: NaiveDate,
end_date: NaiveDate,
want_first_day: bool,
) -> Result<Vec<NaiveDate>, Box<dyn Error>> {
let mut result = Vec::new();
let mut year = start_date.year();
while year <= end_date.year() {
let candidate = if want_first_day {
first_day_of_year(year)?
} else {
last_day_of_year(year)?
};
if candidate >= start_date && candidate <= end_date {
result.push(candidate);
} else if candidate > end_date {
// Optimization: If the candidate date is already past the end_date,
// no subsequent year's candidate will be in range.
break;
}
year = year.checked_add(1).ok_or("Year overflow")?;
}
Ok(result)
}
/// Given a date and a `target_weekday`, returns the date that is the first
/// `target_weekday` on or after the given date.
fn move_to_day_of_week_on_or_after(
date: NaiveDate,
target: Weekday,
) -> Result<NaiveDate, Box<dyn Error>> {
let mut current = date;
while current.weekday() != target {
current = current
.succ_opt()
.ok_or("Date overflow moving to next weekday")?;
}
Ok(current)
}
/// Return the first calendar day of the given (year, month).
fn first_day_of_month(year: i32, month: u32) -> Result<NaiveDate, Box<dyn Error>> {
if !(1..=12).contains(&month) {
return Err(format!("Invalid month: {}", month).into());
}
NaiveDate::from_ymd_opt(year, month, 1)
.ok_or_else(|| format!("Invalid year-month combination: {}-{}", year, month).into())
}
/// Returns the number of days in a given month and year.
fn days_in_month(year: i32, month: u32) -> Result<u32, Box<dyn Error>> {
if !(1..=12).contains(&month) {
return Err(format!("Invalid month: {}", month).into());
}
let (ny, nm) = if month == 12 {
(
year.checked_add(1)
.ok_or("Year overflow calculating next month")?,
1,
)
} else {
(year, month + 1)
};
// Use first_day_of_month which handles ymd creation errors
let first_of_next = first_day_of_month(ny, nm)?;
let last_of_this = first_of_next
.pred_opt()
.ok_or("Date underflow calculating last day of month")?;
Ok(last_of_this.day())
}
/// Return the last calendar day of the given (year, month).
fn last_day_of_month(year: i32, month: u32) -> Result<NaiveDate, Box<dyn Error>> {
// days_in_month now validates month and handles overflow
let last_dom = days_in_month(year, month)?;
NaiveDate::from_ymd_opt(year, month, last_dom)
.ok_or_else(|| format!("Invalid year-month-day: {}-{}-{}", year, month, last_dom).into())
}
/// Converts a month number (1-12) to a quarter number (1-4).
/// Panics if month is invalid (should not happen with valid NaiveDate).
pub fn month_to_quarter(m: u32) -> u32 {
match m {
1..=3 => 1,
4..=6 => 2,
7..=9 => 3,
10..=12 => 4,
_ => panic!("Invalid month: {}", m), // Should only happen with programmer error
}
}
/// Returns the 1st day of the month that starts a given quarter.
fn quarter_start_month(quarter: u32) -> Result<u32, Box<dyn Error>> {
match quarter {
1 => Ok(1), // Jan
2 => Ok(4), // Apr
3 => Ok(7), // Jul
4 => Ok(10), // Oct
_ => Err(format!("invalid quarter: {}", quarter).into()), // Return Err instead of panic
}
}
/// Return the first calendar day in the given (year, quarter).
fn first_day_of_quarter(year: i32, quarter: u32) -> Result<NaiveDate, Box<dyn Error>> {
// Propagate error from quarter_start_month
let month = quarter_start_month(quarter)?;
first_day_of_month(year, month)
}
/// Returns the last day of the month that ends a given quarter.
fn quarter_end_month(quarter: u32) -> Result<u32, Box<dyn Error>> {
match quarter {
1 => Ok(3), // Mar
2 => Ok(6), // Jun
3 => Ok(9), // Sep
4 => Ok(12), // Dec
_ => Err(format!("invalid quarter: {}", quarter).into()), // Return Err instead of panic
}
}
/// Return the last calendar day in the given (year, quarter).
fn last_day_of_quarter(year: i32, quarter: u32) -> Result<NaiveDate, Box<dyn Error>> {
// Propagate error from quarter_end_month
let month = quarter_end_month(quarter)?;
last_day_of_month(year, month)
}
/// Returns the first calendar day (Jan 1st) of the given year.
fn first_day_of_year(year: i32) -> Result<NaiveDate, Box<dyn Error>> {
NaiveDate::from_ymd_opt(year, 1, 1)
.ok_or_else(|| format!("Invalid year for Jan 1st: {}", year).into())
}
/// Returns the last calendar day (Dec 31st) of the given year.
fn last_day_of_year(year: i32) -> Result<NaiveDate, Box<dyn Error>> {
NaiveDate::from_ymd_opt(year, 12, 31)
.ok_or_else(|| format!("Invalid year for Dec 31st: {}", year).into())
}
// --- Generator Helper Functions ---
fn get_first_date_helper(freq: DateFreq) -> fn(i32, u32) -> Result<NaiveDate, Box<dyn Error>> {
if matches!(
freq,
DateFreq::Daily | DateFreq::WeeklyMonday | DateFreq::WeeklyFriday
) {
panic!("Daily, WeeklyMonday, and WeeklyFriday frequencies are not supported here");
}
match freq {
DateFreq::MonthStart => first_day_of_month,
DateFreq::MonthEnd => last_day_of_month,
DateFreq::QuarterStart => first_day_of_quarter,
DateFreq::QuarterEnd => last_day_of_quarter,
DateFreq::YearStart => |year, _| first_day_of_year(year),
DateFreq::YearEnd => |year, _| last_day_of_year(year),
_ => unreachable!(),
}
}
/// Finds the *first* valid date according to the frequency,
/// starting the search *on or after* the given `start_date`.
pub fn find_first_date_on_or_after(
start_date: NaiveDate,
freq: DateFreq,
) -> Result<NaiveDate, Box<dyn Error>> {
match freq {
DateFreq::Daily => Ok(start_date), // The first daily date is the start date itself
DateFreq::WeeklyMonday => move_to_day_of_week_on_or_after(start_date, Weekday::Mon),
DateFreq::WeeklyFriday => move_to_day_of_week_on_or_after(start_date, Weekday::Fri),
DateFreq::MonthStart | DateFreq::MonthEnd => {
// let mut candidate = first_day_of_month(start_date.year(), start_date.month())?;
let get_cand_func = get_first_date_helper(freq);
let mut candidate = get_cand_func(start_date.year(), start_date.month())?;
if candidate < start_date {
let (next_y, next_m) = if start_date.month() == 12 {
(start_date.year().checked_add(1).ok_or("Year overflow")?, 1)
} else {
(start_date.year(), start_date.month() + 1)
};
candidate = get_cand_func(next_y, next_m)?;
}
Ok(candidate)
}
DateFreq::QuarterStart | DateFreq::QuarterEnd => {
let current_q = month_to_quarter(start_date.month());
let get_cand_func = get_first_date_helper(freq);
let mut candidate = get_cand_func(start_date.year(), current_q)?;
if candidate < start_date {
let (next_y, next_q) = if current_q == 4 {
(start_date.year().checked_add(1).ok_or("Year overflow")?, 1)
} else {
(start_date.year(), current_q + 1)
};
candidate = get_cand_func(next_y, next_q)?;
}
Ok(candidate)
}
DateFreq::YearStart | DateFreq::YearEnd => {
let get_cand_func = get_first_date_helper(freq);
let mut candidate = get_cand_func(start_date.year(), 0)?;
if candidate < start_date {
candidate =
get_cand_func(start_date.year().checked_add(1).ok_or("Year overflow")?, 0)?;
}
Ok(candidate)
}
}
}
/// Finds the *next* valid date according to the frequency,
/// given the `current_date` (which is assumed to be a valid date previously generated).
pub fn find_next_date(
current_date: NaiveDate,
freq: DateFreq,
) -> Result<NaiveDate, Box<dyn Error>> {
match freq {
DateFreq::Daily => current_date
.succ_opt()
.ok_or_else(|| "Date overflow finding next daily".into()),
DateFreq::WeeklyMonday | DateFreq::WeeklyFriday => current_date
.checked_add_signed(Duration::days(7))
.ok_or_else(|| "Date overflow adding 7 days".into()),
DateFreq::MonthStart | DateFreq::MonthEnd => {
let get_cand_func = get_first_date_helper(freq);
let (next_y, next_m) = if current_date.month() == 12 {
(
current_date.year().checked_add(1).ok_or("Year overflow")?,
1,
)
} else {
(current_date.year(), current_date.month() + 1)
};
get_cand_func(next_y, next_m)
}
DateFreq::QuarterStart | DateFreq::QuarterEnd => {
let current_q = month_to_quarter(current_date.month());
let get_cand_func = get_first_date_helper(freq);
let (next_y, next_q) = if current_q == 4 {
(
current_date.year().checked_add(1).ok_or("Year overflow")?,
1,
)
} else {
(current_date.year(), current_q + 1)
};
get_cand_func(next_y, next_q)
}
DateFreq::YearStart | DateFreq::YearEnd => {
let get_cand_func = get_first_date_helper(freq);
get_cand_func(
current_date.year().checked_add(1).ok_or("Year overflow")?,
0,
)
}
}
}
// --- Tests ---
#[cfg(test)]
mod tests {
use super::*;
use chrono::{Duration, NaiveDate, Weekday}; // Make sure Duration is imported
// Helper to create a NaiveDate for tests, expecting valid dates.
fn date(year: i32, month: u32, day: u32) -> NaiveDate {
NaiveDate::from_ymd_opt(year, month, day).expect("Invalid date in test setup")
}
// --- DateFreq Tests ---
#[test]
fn test_datefreq_from_str() -> Result<(), Box<dyn Error>> {
assert_eq!(DateFreq::from_str("D")?, DateFreq::Daily);
assert_eq!("D".parse::<DateFreq>()?, DateFreq::Daily);
assert_eq!(DateFreq::from_str("W")?, DateFreq::WeeklyMonday);
assert_eq!(DateFreq::from_str("WS")?, DateFreq::WeeklyMonday);
assert_eq!(DateFreq::from_str("M")?, DateFreq::MonthStart);
assert_eq!(DateFreq::from_str("MS")?, DateFreq::MonthStart);
assert_eq!(DateFreq::from_str("Q")?, DateFreq::QuarterStart);
assert_eq!(DateFreq::from_str("QS")?, DateFreq::QuarterStart);
assert_eq!(DateFreq::from_str("Y")?, DateFreq::YearStart);
assert_eq!(DateFreq::from_str("A")?, DateFreq::YearStart);
assert_eq!(DateFreq::from_str("AS")?, DateFreq::YearStart);
assert_eq!(DateFreq::from_str("YS")?, DateFreq::YearStart);
assert_eq!(DateFreq::from_str("ME")?, DateFreq::MonthEnd);
assert_eq!(DateFreq::from_str("QE")?, DateFreq::QuarterEnd);
assert_eq!(DateFreq::from_str("WF")?, DateFreq::WeeklyFriday);
assert_eq!("WF".parse::<DateFreq>()?, DateFreq::WeeklyFriday);
assert_eq!(DateFreq::from_str("YE")?, DateFreq::YearEnd);
assert_eq!(DateFreq::from_str("AE")?, DateFreq::YearEnd);
assert!(DateFreq::from_str("INVALID").is_err());
assert!("INVALID".parse::<DateFreq>().is_err());
let err = DateFreq::from_str("INVALID").unwrap_err();
assert_eq!(err.to_string(), "Invalid frequency specified: INVALID");
Ok(())
}
#[test]
fn test_datefreq_to_string() {
assert_eq!(DateFreq::Daily.to_string(), "D");
assert_eq!(DateFreq::WeeklyMonday.to_string(), "W");
assert_eq!(DateFreq::MonthStart.to_string(), "M");
assert_eq!(DateFreq::QuarterStart.to_string(), "Q");
assert_eq!(DateFreq::YearStart.to_string(), "Y");
assert_eq!(DateFreq::MonthEnd.to_string(), "ME");
assert_eq!(DateFreq::QuarterEnd.to_string(), "QE");
assert_eq!(DateFreq::WeeklyFriday.to_string(), "WF");
assert_eq!(DateFreq::YearEnd.to_string(), "YE");
}
#[test]
fn test_datefreq_from_string() -> Result<(), Box<dyn Error>> {
assert_eq!(DateFreq::from_string("D".to_string())?, DateFreq::Daily);
assert!(DateFreq::from_string("INVALID".to_string()).is_err());
Ok(())
}
#[test]
fn test_datefreq_agg_type() {
assert_eq!(DateFreq::Daily.agg_type(), AggregationType::Start);
assert_eq!(DateFreq::WeeklyMonday.agg_type(), AggregationType::Start);
assert_eq!(DateFreq::MonthStart.agg_type(), AggregationType::Start);
assert_eq!(DateFreq::QuarterStart.agg_type(), AggregationType::Start);
assert_eq!(DateFreq::YearStart.agg_type(), AggregationType::Start);
assert_eq!(DateFreq::WeeklyFriday.agg_type(), AggregationType::End);
assert_eq!(DateFreq::MonthEnd.agg_type(), AggregationType::End);
assert_eq!(DateFreq::QuarterEnd.agg_type(), AggregationType::End);
assert_eq!(DateFreq::YearEnd.agg_type(), AggregationType::End);
}
// --- DatesList Property Tests ---
#[test]
fn test_dates_list_properties_new() -> Result<(), Box<dyn Error>> {
let start_str = "2023-01-01".to_string();
let end_str = "2023-12-31".to_string();
let freq = DateFreq::QuarterEnd;
let dates_list = DatesList::new(start_str.clone(), end_str.clone(), freq);
assert_eq!(dates_list.start_date_str(), start_str);
assert_eq!(dates_list.end_date_str(), end_str);
assert_eq!(dates_list.freq(), freq);
assert_eq!(dates_list.freq_str(), "QE");
assert_eq!(dates_list.start_date()?, date(2023, 1, 1));
assert_eq!(dates_list.end_date()?, date(2023, 12, 31));
Ok(())
}
#[test]
fn test_dates_list_properties_from_n_periods() -> Result<(), Box<dyn Error>> {
let start_str = "2023-01-01".to_string(); // Sunday
let freq = DateFreq::Daily;
let n_periods = 5; // Expect: Jan 1, 2, 3, 4, 5
let dates_list = DatesList::from_n_periods(start_str.clone(), freq, n_periods)?;
assert_eq!(dates_list.start_date_str(), start_str);
assert_eq!(dates_list.end_date_str(), "2023-01-05");
assert_eq!(dates_list.freq(), freq);
assert_eq!(dates_list.freq_str(), "D");
assert_eq!(dates_list.start_date()?, date(2023, 1, 1));
assert_eq!(dates_list.end_date()?, date(2023, 1, 5));
assert_eq!(
dates_list.list()?,
vec![
date(2023, 1, 1),
date(2023, 1, 2),
date(2023, 1, 3),
date(2023, 1, 4),
date(2023, 1, 5)
]
);
assert_eq!(dates_list.count()?, 5);
Ok(())
}
#[test]
fn test_dates_list_from_n_periods_zero_periods() {
let start_str = "2023-01-01".to_string();
let freq = DateFreq::Daily;
let n_periods = 0;
let result = DatesList::from_n_periods(start_str.clone(), freq, n_periods);
assert!(result.is_err());
assert_eq!(
result.unwrap_err().to_string(),
"n_periods must be greater than 0"
);
}
// test_dates_list_from_n_periods_fail_get_last removed as it was flawed
#[test]
fn test_dates_list_from_n_periods_invalid_start_date() {
let start_str = "invalid-date".to_string();
let freq = DateFreq::Daily;
let n_periods = 5;
let result = DatesList::from_n_periods(start_str.clone(), freq, n_periods);
assert!(result.is_err());
assert!(
result
.unwrap_err()
.to_string()
.contains("input contains invalid characters") // Error from NaiveDate::parse_from_str
);
}
#[test]
fn test_dates_list_invalid_date_string_new() {
let list_start_invalid = DatesList::new(
"invalid-date".to_string(),
"2023-12-31".to_string(),
DateFreq::Daily,
);
assert!(list_start_invalid.list().is_err());
assert!(list_start_invalid.count().is_err());
assert!(list_start_invalid.groups().is_err());
assert!(list_start_invalid.start_date().is_err());
assert!(list_start_invalid.end_date().is_ok()); // End date is valid
let list_end_invalid = DatesList::new(
"2023-01-01".to_string(),
"invalid-date".to_string(),
DateFreq::Daily,
);
assert!(list_end_invalid.list().is_err());
assert!(list_end_invalid.count().is_err());
assert!(list_end_invalid.groups().is_err());
assert!(list_end_invalid.start_date().is_ok()); // Start date is valid
assert!(list_end_invalid.end_date().is_err());
}
// --- DatesList Core Logic Tests (list, count) ---
#[test]
fn test_dates_list_daily_list() -> Result<(), Box<dyn Error>> {
let dates_list = DatesList::new(
"2023-11-01".to_string(), // Wed
"2023-11-05".to_string(), // Sun
DateFreq::Daily,
);
let list = dates_list.list()?;
assert_eq!(list.len(), 5);
assert_eq!(
list,
vec![
date(2023, 11, 1),
date(2023, 11, 2),
date(2023, 11, 3),
date(2023, 11, 4),
date(2023, 11, 5)
]
);
assert_eq!(dates_list.count()?, 5);
Ok(())
}
#[test]
fn test_dates_list_weekly_monday_list() -> Result<(), Box<dyn Error>> {
let dates_list = DatesList::new(
"2023-10-30".to_string(), // Mon
"2023-11-13".to_string(), // Mon
DateFreq::WeeklyMonday,
);
let list = dates_list.list()?;
// Mondays in range: Oct 30, Nov 6, Nov 13
assert_eq!(list.len(), 3);
assert_eq!(
list,
vec![date(2023, 10, 30), date(2023, 11, 6), date(2023, 11, 13)]
);
assert_eq!(dates_list.count()?, 3);
Ok(())
}
#[test]
fn test_dates_list_weekly_friday_list() -> Result<(), Box<dyn Error>> {
let dates_list = DatesList::new(
"2023-11-01".to_string(), // Wed
"2023-11-17".to_string(), // Fri
DateFreq::WeeklyFriday,
);
let list = dates_list.list()?;
// Fridays in range: Nov 3, Nov 10, Nov 17
assert_eq!(list.len(), 3);
assert_eq!(
list,
vec![date(2023, 11, 3), date(2023, 11, 10), date(2023, 11, 17)]
);
assert_eq!(dates_list.count()?, 3);
Ok(())
}
#[test]
fn test_dates_list_month_start_list() -> Result<(), Box<dyn Error>> {
let dates_list = DatesList::new(
"2023-11-15".to_string(), // Mid-Nov
"2024-02-01".to_string(), // Feb 1st
DateFreq::MonthStart,
);
let list = dates_list.list()?;
// Month starts >= Nov 15 and <= Feb 1: Dec 1, Jan 1, Feb 1
assert_eq!(list.len(), 3);
assert_eq!(
list,
vec![date(2023, 12, 1), date(2024, 1, 1), date(2024, 2, 1)]
);
assert_eq!(dates_list.count()?, 3);
Ok(())
}
#[test]
fn test_dates_list_month_end_list_leap() -> Result<(), Box<dyn Error>> {
let dates_list = DatesList::new(
"2024-01-15".to_string(), // Mid-Jan
"2024-03-31".to_string(), // Mar 31st
DateFreq::MonthEnd,
);
let list = dates_list.list()?;
// Month ends >= Jan 15 and <= Mar 31: Jan 31, Feb 29 (leap), Mar 31
assert_eq!(list.len(), 3);
assert_eq!(
list,
vec![date(2024, 1, 31), date(2024, 2, 29), date(2024, 3, 31)]
);
assert_eq!(dates_list.count()?, 3);
Ok(())
}
#[test]
fn test_dates_list_quarter_start_list() -> Result<(), Box<dyn Error>> {
let dates_list = DatesList::new(
"2023-08-01".to_string(), // Mid Q3
"2024-04-01".to_string(), // Start Q2
DateFreq::QuarterStart,
);
let list = dates_list.list()?;
// Quarter starts >= Aug 1 '23 and <= Apr 1 '24: Oct 1 '23, Jan 1 '24, Apr 1 '24
assert_eq!(list.len(), 3);
assert_eq!(
list,
vec![date(2023, 10, 1), date(2024, 1, 1), date(2024, 4, 1)]
);
assert_eq!(dates_list.count()?, 3);
Ok(())
}
#[test]
fn test_dates_list_quarter_end_list() -> Result<(), Box<dyn Error>> {
let dates_list = DatesList::new(
"2023-03-31".to_string(), // End Q1
"2023-12-31".to_string(), // End Q4
DateFreq::QuarterEnd,
);
let list = dates_list.list()?;
// Quarter ends >= Mar 31 and <= Dec 31: Mar 31, Jun 30, Sep 30, Dec 31
assert_eq!(list.len(), 4);
assert_eq!(
list,
vec![
date(2023, 3, 31),
date(2023, 6, 30),
date(2023, 9, 30),
date(2023, 12, 31)
]
);
assert_eq!(dates_list.count()?, 4);
Ok(())
}
#[test]
fn test_dates_list_year_start_list() -> Result<(), Box<dyn Error>> {
let dates_list = DatesList::new(
"2023-06-01".to_string(), // Mid 2023
"2025-01-01".to_string(), // Start 2025
DateFreq::YearStart,
);
let list = dates_list.list()?;
// Year starts >= Jun 1 '23 and <= Jan 1 '25: Jan 1 '24, Jan 1 '25
assert_eq!(list.len(), 2);
assert_eq!(list, vec![date(2024, 1, 1), date(2025, 1, 1)]);
assert_eq!(dates_list.count()?, 2);
Ok(())
}
#[test]
fn test_dates_list_year_end_list() -> Result<(), Box<dyn Error>> {
let dates_list = DatesList::new(
"2022-01-01".to_string(), // Start 2022
"2024-03-31".to_string(), // Q1 2024
DateFreq::YearEnd,
);
let list = dates_list.list()?;
// Year ends >= Jan 1 '22 and <= Mar 31 '24: Dec 31 '22, Dec 31 '23
assert_eq!(list.len(), 2);
assert_eq!(list, vec![date(2022, 12, 31), date(2023, 12, 31)]);
assert_eq!(dates_list.count()?, 2);
Ok(())
}
#[test]
fn test_dates_list_empty_range_list() -> Result<(), Box<dyn Error>> {
let dates_list = DatesList::new(
"2023-12-31".to_string(),
"2023-01-01".to_string(), // End date before start date
DateFreq::Daily,
);
let list = dates_list.list()?;
assert!(list.is_empty());
assert_eq!(dates_list.count()?, 0);
Ok(())
}
// --- Tests for groups() method ---
#[test]
fn test_dates_list_groups_monthly_end() -> Result<(), Box<dyn Error>> {
let dates_list = DatesList::new(
"2023-10-15".to_string(), // Mid-Oct
"2024-01-15".to_string(), // Mid-Jan
DateFreq::MonthEnd,
);
let groups = dates_list.groups()?;
// Month Ends >= Oct 15 '23 and <= Jan 15 '24: Oct 31, Nov 30, Dec 31
assert_eq!(groups.len(), 3);
// Key order: Monthly(2023, 10), Monthly(2023, 11), Monthly(2023, 12)
assert_eq!(groups[0], vec![date(2023, 10, 31)]);
assert_eq!(groups[1], vec![date(2023, 11, 30)]);
assert_eq!(groups[2], vec![date(2023, 12, 31)]);
Ok(())
}
#[test]
fn test_dates_list_groups_daily() -> Result<(), Box<dyn Error>> {
let dates_list = DatesList::new(
"2023-11-01".to_string(), // Wed
"2023-11-03".to_string(), // Fri
DateFreq::Daily,
);
let groups = dates_list.groups()?;
// Dates: Nov 1, Nov 2, Nov 3. Each gets own group.
assert_eq!(groups.len(), 3);
// Key order: Daily(2023-11-01), Daily(2023-11-02), Daily(2023-11-03)
assert_eq!(groups[0], vec![date(2023, 11, 1)]);
assert_eq!(groups[1], vec![date(2023, 11, 2)]);
assert_eq!(groups[2], vec![date(2023, 11, 3)]);
Ok(())
}
#[test]
fn test_dates_list_groups_weekly_friday() -> Result<(), Box<dyn Error>> {
let dates_list = DatesList::new(
"2023-11-01".to_string(), // Wed (ISO Week 44)
"2023-11-15".to_string(), // Wed (ISO Week 46)
DateFreq::WeeklyFriday,
);
let groups = dates_list.groups()?;
// Fridays in range: Nov 3 (W44), Nov 10 (W45)
assert_eq!(groups.len(), 2);
// Key order: Weekly(2023, 44), Weekly(2023, 45)
assert_eq!(groups[0], vec![date(2023, 11, 3)]);
assert_eq!(groups[1], vec![date(2023, 11, 10)]);
Ok(())
}
#[test]
fn test_dates_list_groups_quarterly_start() -> Result<(), Box<dyn Error>> {
let dates_list = DatesList::new(
"2023-08-01".to_string(), // Start Q3
"2024-05-01".to_string(), // Start Q2
DateFreq::QuarterStart,
);
let groups = dates_list.groups()?;
// Quarter starts >= Aug 1 '23 and <= May 1 '24: Oct 1 '23, Jan 1 '24, Apr 1 '24
assert_eq!(groups.len(), 3);
// Key order: Quarterly(2023, 4), Quarterly(2024, 1), Quarterly(2024, 2)
assert_eq!(groups[0], vec![date(2023, 10, 1)]);
assert_eq!(groups[1], vec![date(2024, 1, 1)]);
assert_eq!(groups[2], vec![date(2024, 4, 1)]);
Ok(())
}
#[test]
fn test_dates_list_groups_yearly_end() -> Result<(), Box<dyn Error>> {
let dates_list = DatesList::new(
"2022-01-01".to_string(), // Start 2022
"2024-03-31".to_string(), // Q1 2024
DateFreq::YearEnd,
);
let groups = dates_list.groups()?;
// Year ends >= Jan 1 '22 and <= Mar 31 '24: Dec 31 '22, Dec 31 '23
assert_eq!(groups.len(), 2);
// Key order: Yearly(2022), Yearly(2023)
assert_eq!(groups[0], vec![date(2022, 12, 31)]);
assert_eq!(groups[1], vec![date(2023, 12, 31)]);
Ok(())
}
#[test]
fn test_dates_list_groups_empty_range() -> Result<(), Box<dyn Error>> {
let dates_list = DatesList::new(
"2023-12-31".to_string(),
"2023-01-01".to_string(), // End < Start
DateFreq::Daily,
);
let groups = dates_list.groups()?;
assert!(groups.is_empty());
Ok(())
}
#[test]
fn test_move_to_day_of_week_on_or_after() -> Result<(), Box<dyn Error>> {
assert_eq!(
move_to_day_of_week_on_or_after(date(2023, 11, 6), Weekday::Mon)?,
date(2023, 11, 6)
);
assert_eq!(
move_to_day_of_week_on_or_after(date(2023, 11, 8), Weekday::Fri)?,
date(2023, 11, 10)
);
assert_eq!(
move_to_day_of_week_on_or_after(date(2023, 11, 11), Weekday::Mon)?, // Sat -> Mon
date(2023, 11, 13)
);
assert_eq!(
move_to_day_of_week_on_or_after(date(2023, 11, 10), Weekday::Mon)?, // Fri -> Mon
date(2023, 11, 13)
);
// Test near max date (ensure it doesn't panic easily, though overflow is possible)
// MAX - 7 days guarantees we have room to move forward
let near_max = NaiveDate::MAX - Duration::days(7);
assert!(move_to_day_of_week_on_or_after(near_max, Weekday::Sun).is_ok());
// Test overflow case - starting at MAX, moving forward fails if MAX is not target
if NaiveDate::MAX.weekday() != Weekday::Sun {
assert!(move_to_day_of_week_on_or_after(NaiveDate::MAX, Weekday::Sun).is_err());
} else {
// If MAX is the target, it should succeed
assert!(move_to_day_of_week_on_or_after(NaiveDate::MAX, Weekday::Sun).is_ok());
// And trying to move *past* it should fail
let day_before = NaiveDate::MAX - Duration::days(1);
let target_day_after = NaiveDate::MAX.weekday().succ(); // Day after MAX's weekday
assert!(move_to_day_of_week_on_or_after(day_before, target_day_after).is_err());
// Moving past MAX fails
}
Ok(())
}
#[test]
fn test_first_day_of_month() -> Result<(), Box<dyn Error>> {
assert_eq!(first_day_of_month(2023, 11)?, date(2023, 11, 1));
assert_eq!(first_day_of_month(2024, 2)?, date(2024, 2, 1));
assert!(first_day_of_month(2023, 0).is_err()); // Invalid month 0
assert!(first_day_of_month(2023, 13).is_err()); // Invalid month 13
Ok(())
}
#[test]
fn test_days_in_month() -> Result<(), Box<dyn Error>> {
assert_eq!(days_in_month(2023, 1)?, 31);
assert_eq!(days_in_month(2023, 2)?, 28);
// Leap
assert_eq!(days_in_month(2024, 2)?, 29);
assert_eq!(days_in_month(2023, 4)?, 30);
assert_eq!(days_in_month(2023, 12)?, 31);
// Invalid month 0
assert!(days_in_month(2023, 0).is_err());
// Invalid month 13
// Test near max date year overflow - Use MAX.year()
assert!(days_in_month(2023, 13).is_err());
assert!(days_in_month(NaiveDate::MAX.year(), 12).is_err());
Ok(())
}
#[test]
fn test_last_day_of_month() -> Result<(), Box<dyn Error>> {
assert_eq!(last_day_of_month(2023, 11)?, date(2023, 11, 30));
assert_eq!(last_day_of_month(2024, 2)?, date(2024, 2, 29)); // Leap
assert_eq!(last_day_of_month(2023, 12)?, date(2023, 12, 31));
// Invalid month 0
assert!(last_day_of_month(2023, 0).is_err());
// Invalid month 13
// Test near max date year overflow - use MAX.year()
assert!(last_day_of_month(2023, 13).is_err());
assert!(last_day_of_month(NaiveDate::MAX.year(), 12).is_err());
Ok(())
}
#[test]
fn test_month_to_quarter() {
assert_eq!(month_to_quarter(1), 1);
assert_eq!(month_to_quarter(3), 1);
assert_eq!(month_to_quarter(4), 2);
assert_eq!(month_to_quarter(6), 2);
assert_eq!(month_to_quarter(7), 3);
assert_eq!(month_to_quarter(9), 3);
assert_eq!(month_to_quarter(10), 4);
assert_eq!(month_to_quarter(12), 4);
}
#[test]
#[should_panic(expected = "Invalid month: 0")]
fn test_month_to_quarter_invalid_low() {
month_to_quarter(0);
}
#[test]
#[should_panic(expected = "Invalid month: 13")]
fn test_month_to_quarter_invalid_high() {
month_to_quarter(13);
}
#[test]
fn test_quarter_start_month() {
assert_eq!(quarter_start_month(1).unwrap(), 1);
assert_eq!(quarter_start_month(2).unwrap(), 4);
assert_eq!(quarter_start_month(3).unwrap(), 7);
assert_eq!(quarter_start_month(4).unwrap(), 10);
assert!(quarter_start_month(0).is_err());
assert!(quarter_start_month(5).is_err());
}
#[test]
fn test_first_day_of_quarter() -> Result<(), Box<dyn Error>> {
assert_eq!(first_day_of_quarter(2023, 1)?, date(2023, 1, 1));
assert_eq!(first_day_of_quarter(2023, 2)?, date(2023, 4, 1));
assert_eq!(first_day_of_quarter(2023, 3)?, date(2023, 7, 1));
assert_eq!(first_day_of_quarter(2023, 4)?, date(2023, 10, 1));
// Invalid quarter
assert!(first_day_of_quarter(2023, 5).is_err());
Ok(())
}
#[test]
fn test_quarter_end_month() {
assert_eq!(quarter_end_month(1).unwrap(), 3);
assert_eq!(quarter_end_month(2).unwrap(), 6);
assert_eq!(quarter_end_month(3).unwrap(), 9);
assert_eq!(quarter_end_month(4).unwrap(), 12);
assert!(quarter_end_month(0).is_err());
assert!(quarter_end_month(5).is_err());
}
#[test]
fn test_last_day_of_quarter() -> Result<(), Box<dyn Error>> {
assert_eq!(last_day_of_quarter(2023, 1)?, date(2023, 3, 31));
assert_eq!(last_day_of_quarter(2023, 2)?, date(2023, 6, 30));
assert_eq!(last_day_of_quarter(2023, 3)?, date(2023, 9, 30));
assert_eq!(last_day_of_quarter(2023, 4)?, date(2023, 12, 31));
// Leap year doesn't affect March end
assert_eq!(last_day_of_quarter(2024, 1)?, date(2024, 3, 31));
// Invalid quarter
// Test overflow propagation - use MAX.year()
assert!(last_day_of_quarter(2023, 5).is_err());
assert!(last_day_of_quarter(NaiveDate::MAX.year(), 4).is_err());
Ok(())
}
#[test]
fn test_first_day_of_year() -> Result<(), Box<dyn Error>> {
assert_eq!(first_day_of_year(2023)?, date(2023, 1, 1));
assert_eq!(first_day_of_year(2024)?, date(2024, 1, 1));
// Test MAX year - should be ok for Jan 1
assert!(first_day_of_year(NaiveDate::MAX.year()).is_ok());
Ok(())
}
#[test]
fn test_last_day_of_year() -> Result<(), Box<dyn Error>> {
assert_eq!(last_day_of_year(2023)?, date(2023, 12, 31));
// Leap year doesn't affect Dec 31st existence
// Test MAX year - should be okay since MAX is Dec 31
assert_eq!(last_day_of_year(2024)?, date(2024, 12, 31));
assert_eq!(last_day_of_year(NaiveDate::MAX.year())?, NaiveDate::MAX);
Ok(())
}
#[test]
fn test_find_first_date_on_or_after() -> Result<(), Box<dyn Error>> {
// Daily
assert_eq!(
find_first_date_on_or_after(date(2023, 11, 8), DateFreq::Daily)?,
date(2023, 11, 8)
);
// Sat -> Sat
assert_eq!(
find_first_date_on_or_after(date(2023, 11, 11), DateFreq::Daily)?,
date(2023, 11, 11)
);
// Weekly Mon
assert_eq!(
find_first_date_on_or_after(date(2023, 11, 8), DateFreq::WeeklyMonday)?,
date(2023, 11, 13)
);
assert_eq!(
find_first_date_on_or_after(date(2023, 11, 13), DateFreq::WeeklyMonday)?,
date(2023, 11, 13)
);
// Sun -> Mon
assert_eq!(
find_first_date_on_or_after(date(2023, 11, 12), DateFreq::WeeklyMonday)?,
date(2023, 11, 13)
);
// Weekly Fri
assert_eq!(
find_first_date_on_or_after(date(2023, 11, 8), DateFreq::WeeklyFriday)?,
date(2023, 11, 10)
);
assert_eq!(
find_first_date_on_or_after(date(2023, 11, 10), DateFreq::WeeklyFriday)?,
date(2023, 11, 10)
);
assert_eq!(
find_first_date_on_or_after(date(2023, 11, 11), DateFreq::WeeklyFriday)?,
date(2023, 11, 17)
); // Sat -> Next Fri
// Month Start
assert_eq!(
find_first_date_on_or_after(date(2023, 11, 1), DateFreq::MonthStart)?,
date(2023, 11, 1)
);
assert_eq!(
find_first_date_on_or_after(date(2023, 10, 15), DateFreq::MonthStart)?,
date(2023, 11, 1)
);
assert_eq!(
find_first_date_on_or_after(date(2023, 12, 15), DateFreq::MonthStart)?,
date(2024, 1, 1)
);
// Oct 1 -> Oct 1
assert_eq!(
find_first_date_on_or_after(date(2023, 10, 1), DateFreq::MonthStart)?,
date(2023, 10, 1)
);
// Month End
assert_eq!(
find_first_date_on_or_after(date(2023, 11, 30), DateFreq::MonthEnd)?,
date(2023, 11, 30)
);
assert_eq!(
find_first_date_on_or_after(date(2023, 11, 15), DateFreq::MonthEnd)?,
date(2023, 11, 30)
);
// Dec 31 -> Dec 31
assert_eq!(
find_first_date_on_or_after(date(2023, 12, 31), DateFreq::MonthEnd)?,
date(2023, 12, 31)
);
// Mid Feb (Leap) -> Feb 29
assert_eq!(
find_first_date_on_or_after(date(2024, 2, 15), DateFreq::MonthEnd)?,
date(2024, 2, 29)
);
// Feb 29 -> Feb 29
assert_eq!(
find_first_date_on_or_after(date(2024, 2, 29), DateFreq::MonthEnd)?,
date(2024, 2, 29)
);
// Quarter Start
assert_eq!(
find_first_date_on_or_after(date(2023, 10, 1), DateFreq::QuarterStart)?,
date(2023, 10, 1)
);
assert_eq!(
find_first_date_on_or_after(date(2023, 8, 15), DateFreq::QuarterStart)?,
date(2023, 10, 1)
);
assert_eq!(
find_first_date_on_or_after(date(2023, 11, 15), DateFreq::QuarterStart)?,
date(2024, 1, 1)
);
assert_eq!(
find_first_date_on_or_after(date(2023, 1, 1), DateFreq::QuarterStart)?,
date(2023, 1, 1)
);
// Quarter End
assert_eq!(
find_first_date_on_or_after(date(2023, 9, 30), DateFreq::QuarterEnd)?,
date(2023, 9, 30)
);
assert_eq!(
find_first_date_on_or_after(date(2023, 8, 15), DateFreq::QuarterEnd)?,
date(2023, 9, 30)
);
assert_eq!(
find_first_date_on_or_after(date(2023, 10, 15), DateFreq::QuarterEnd)?,
date(2023, 12, 31)
);
assert_eq!(
find_first_date_on_or_after(date(2023, 12, 31), DateFreq::QuarterEnd)?,
date(2023, 12, 31)
);
// Year Start
assert_eq!(
find_first_date_on_or_after(date(2024, 1, 1), DateFreq::YearStart)?,
date(2024, 1, 1)
);
assert_eq!(
find_first_date_on_or_after(date(2023, 6, 15), DateFreq::YearStart)?,
date(2024, 1, 1)
);
assert_eq!(
find_first_date_on_or_after(date(2023, 1, 1), DateFreq::YearStart)?,
date(2023, 1, 1)
);
// Year End
assert_eq!(
find_first_date_on_or_after(date(2023, 12, 31), DateFreq::YearEnd)?,
date(2023, 12, 31)
);
assert_eq!(
find_first_date_on_or_after(date(2023, 6, 15), DateFreq::YearEnd)?,
date(2023, 12, 31)
);
assert_eq!(
find_first_date_on_or_after(date(2022, 12, 31), DateFreq::YearEnd)?,
date(2022, 12, 31)
);
// --- Test Overflow Cases near MAX ---
assert!(find_first_date_on_or_after(NaiveDate::MAX, DateFreq::Daily).is_ok()); // Daily starting at MAX is MAX
// Weekly: depends if MAX is the target day. If not, succ() fails.
if NaiveDate::MAX.weekday() != Weekday::Mon {
assert!(find_first_date_on_or_after(NaiveDate::MAX, DateFreq::WeeklyMonday).is_err());
} else {
assert!(find_first_date_on_or_after(NaiveDate::MAX, DateFreq::WeeklyMonday).is_ok());
}
// Month Start: MAX is Dec 31. find_first for MonthStart at MAX tries month=12, candidate=Dec 1. candidate < MAX is true.
// Tries next month: Jan (MAX_YEAR+1), which fails in checked_add.
assert!(find_first_date_on_or_after(NaiveDate::MAX, DateFreq::MonthStart).is_err());
// Month End: MAX is Dec 31. find_first for MonthEnd at MAX tries month=12, calls last_day_of_month(MAX_YEAR, 12).
// last_day_of_month -> days_in_month -> first_day_of_month(MAX_YEAR+1, 1) -> fails.
assert!(find_first_date_on_or_after(NaiveDate::MAX, DateFreq::MonthEnd).is_err());
// Quarter Start: MAX is Dec 31 (Q4). Tries Q4 start (Oct 1). candidate < MAX is true. Tries next Q (Q1 MAX+1), fails.
assert!(find_first_date_on_or_after(NaiveDate::MAX, DateFreq::QuarterStart).is_err());
// Quarter End: MAX is Dec 31 (Q4). Tries Q4 end (Dec 31). Calls last_day_of_quarter(MAX_YEAR, 4).
// last_day_of_quarter -> last_day_of_month(MAX_YEAR, 12) -> fails.
assert!(find_first_date_on_or_after(NaiveDate::MAX, DateFreq::QuarterEnd).is_err());
// Year Start: MAX is Dec 31. Tries YearStart(MAX_YEAR) (Jan 1). candidate < MAX is true. Tries next year (MAX_YEAR+1), fails.
assert!(find_first_date_on_or_after(NaiveDate::MAX, DateFreq::YearStart).is_err());
// Year End: MAX is Dec 31. Tries YearEnd(MAX_YEAR). Calls last_day_of_year(MAX_YEAR). Returns Ok(MAX). candidate < MAX is false. Returns Ok(MAX).
assert!(find_first_date_on_or_after(NaiveDate::MAX, DateFreq::YearEnd).is_ok());
Ok(())
}
#[test]
fn test_find_next_date() -> Result<(), Box<dyn Error>> {
// Daily
assert_eq!(
find_next_date(date(2023, 11, 8), DateFreq::Daily)?,
date(2023, 11, 9)
);
assert_eq!(
find_next_date(date(2023, 11, 10), DateFreq::Daily)?,
date(2023, 11, 11)
); // Fri -> Sat
// Weekly Mon
assert_eq!(
find_next_date(date(2023, 11, 13), DateFreq::WeeklyMonday)?,
date(2023, 11, 20)
);
// Weekly Fri
assert_eq!(
find_next_date(date(2023, 11, 10), DateFreq::WeeklyFriday)?,
date(2023, 11, 17)
);
// Month Start
assert_eq!(
find_next_date(date(2023, 11, 1), DateFreq::MonthStart)?,
date(2023, 12, 1)
);
assert_eq!(
find_next_date(date(2023, 12, 1), DateFreq::MonthStart)?,
date(2024, 1, 1)
);
// Month End
assert_eq!(
find_next_date(date(2023, 10, 31), DateFreq::MonthEnd)?,
date(2023, 11, 30)
);
assert_eq!(
find_next_date(date(2024, 1, 31), DateFreq::MonthEnd)?,
date(2024, 2, 29)
); // Jan -> Feb (Leap)
assert_eq!(
find_next_date(date(2024, 2, 29), DateFreq::MonthEnd)?,
date(2024, 3, 31)
); // Feb -> Mar
// Quarter Start
assert_eq!(
find_next_date(date(2023, 10, 1), DateFreq::QuarterStart)?,
date(2024, 1, 1)
);
assert_eq!(
find_next_date(date(2024, 1, 1), DateFreq::QuarterStart)?,
date(2024, 4, 1)
);
// Quarter End
assert_eq!(
find_next_date(date(2023, 9, 30), DateFreq::QuarterEnd)?,
date(2023, 12, 31)
);
assert_eq!(
find_next_date(date(2023, 12, 31), DateFreq::QuarterEnd)?,
date(2024, 3, 31)
);
// Year Start
assert_eq!(
find_next_date(date(2023, 1, 1), DateFreq::YearStart)?,
date(2024, 1, 1)
);
assert_eq!(
find_next_date(date(2024, 1, 1), DateFreq::YearStart)?,
date(2025, 1, 1)
);
// Year End
assert_eq!(
find_next_date(date(2022, 12, 31), DateFreq::YearEnd)?,
date(2023, 12, 31)
);
assert_eq!(
find_next_date(date(2023, 12, 31), DateFreq::YearEnd)?,
date(2024, 12, 31)
);
// --- Test Overflow Cases near MAX ---
assert!(find_next_date(NaiveDate::MAX, DateFreq::Daily).is_err());
assert!(
find_next_date(NaiveDate::MAX - Duration::days(6), DateFreq::WeeklyMonday).is_err()
);
// Test finding next month start after Dec MAX_YEAR -> Jan (MAX_YEAR+1) (fail)
assert!(find_next_date(date(NaiveDate::MAX.year(), 12, 1), DateFreq::MonthStart).is_err());
// Test finding next month end after Nov MAX_YEAR -> Dec MAX_YEAR (fails because last_day_of_month(MAX, 12) fails)
let nov_end_max_year = last_day_of_month(NaiveDate::MAX.year(), 11)?;
assert!(find_next_date(nov_end_max_year, DateFreq::MonthEnd).is_err());
// Test finding next month end after Dec MAX_YEAR -> Jan (MAX_YEAR+1) (fail)
// The call last_day_of_month(MAX_YEAR + 1, 1) fails
assert!(find_next_date(NaiveDate::MAX, DateFreq::MonthEnd).is_err());
// Test finding next quarter start after Q4 MAX_YEAR -> Q1 (MAX_YEAR+1) (fail)
assert!(find_next_date(
first_day_of_quarter(NaiveDate::MAX.year(), 4)?,
DateFreq::QuarterStart
)
.is_err());
// Test finding next quarter end after Q3 MAX_YEAR -> Q4 MAX_YEAR (fails because last_day_of_quarter(MAX, 4) fails)
let q3_end_max_year = last_day_of_quarter(NaiveDate::MAX.year(), 3)?;
assert!(find_next_date(q3_end_max_year, DateFreq::QuarterEnd).is_err());
// Test finding next quarter end after Q4 MAX_YEAR -> Q1 (MAX_YEAR+1) (fail)
// The call last_day_of_quarter(MAX_YEAR + 1, 1) fails
assert!(find_next_date(NaiveDate::MAX, DateFreq::QuarterEnd).is_err());
// Test finding next year start after Jan 1 MAX_YEAR -> Jan 1 (MAX_YEAR+1) (fail)
assert!(find_next_date(
first_day_of_year(NaiveDate::MAX.year())?,
DateFreq::YearStart
)
.is_err());
// Test finding next year end after Dec 31 (MAX_YEAR-1) -> Dec 31 MAX_YEAR (ok)
assert!(find_next_date(
last_day_of_year(NaiveDate::MAX.year() - 1)?,
DateFreq::YearEnd
)
.is_ok());
// Test finding next year end after Dec 31 MAX_YEAR -> Dec 31 (MAX_YEAR+1) (fail)
assert!(
find_next_date(last_day_of_year(NaiveDate::MAX.year())?, DateFreq::YearEnd).is_err()
); // Fails calculating MAX_YEAR+1
Ok(())
}
// --- Tests for DatesGenerator ---
#[test]
fn test_generator_new_zero_periods() -> Result<(), Box<dyn Error>> {
let start_date = date(2023, 1, 1);
let freq = DateFreq::Daily;
let n_periods = 0;
let mut generator = DatesGenerator::new(start_date, freq, n_periods)?;
assert_eq!(generator.next(), None); // Immediately exhausted
Ok(())
}
#[test]
fn test_generator_new_fail_find_first() -> Result<(), Box<dyn Error>> {
let start_date = NaiveDate::MAX;
// Use a frequency that requires finding the *next* day if MAX isn't the target.
let freq = DateFreq::WeeklyMonday;
let n_periods = 1;
let result = DatesGenerator::new(start_date, freq, n_periods);
// This fails if MAX is not a Monday, because find_first tries MAX.succ_opt()
if NaiveDate::MAX.weekday() != Weekday::Mon {
assert!(result.is_err());
} else {
// If MAX *is* a Monday, new() succeeds.
assert!(result.is_ok());
}
Ok(())
}
#[test]
fn test_generator_daily() -> Result<(), Box<dyn Error>> {
let start_date = date(2023, 11, 10); // Fri
let freq = DateFreq::Daily;
let n_periods = 4;
let mut generator = DatesGenerator::new(start_date, freq, n_periods)?;
assert_eq!(generator.next(), Some(date(2023, 11, 10))); // Fri
assert_eq!(generator.next(), Some(date(2023, 11, 11))); // Sat
assert_eq!(generator.next(), Some(date(2023, 11, 12))); // Sun
assert_eq!(generator.next(), Some(date(2023, 11, 13))); // Mon
assert_eq!(generator.next(), None); // Exhausted
// Test collecting
let generator_collect = DatesGenerator::new(start_date, freq, n_periods)?;
assert_eq!(
generator_collect.collect::<Vec<_>>(),
vec![
date(2023, 11, 10),
date(2023, 11, 11),
date(2023, 11, 12),
date(2023, 11, 13)
]
);
Ok(())
}
#[test]
fn test_generator_weekly_monday() -> Result<(), Box<dyn Error>> {
let start_date = date(2023, 11, 8); // Wed
let freq = DateFreq::WeeklyMonday;
let n_periods = 3;
let mut generator = DatesGenerator::new(start_date, freq, n_periods)?;
assert_eq!(generator.next(), Some(date(2023, 11, 13)));
assert_eq!(generator.next(), Some(date(2023, 11, 20)));
assert_eq!(generator.next(), Some(date(2023, 11, 27)));
assert_eq!(generator.next(), None);
Ok(())
}
#[test]
fn test_generator_weekly_friday() -> Result<(), Box<dyn Error>> {
let start_date = date(2023, 11, 11); // Sat
let freq = DateFreq::WeeklyFriday;
let n_periods = 3;
let mut generator = DatesGenerator::new(start_date, freq, n_periods)?;
assert_eq!(generator.next(), Some(date(2023, 11, 17)));
assert_eq!(generator.next(), Some(date(2023, 11, 24)));
assert_eq!(generator.next(), Some(date(2023, 12, 1)));
assert_eq!(generator.next(), None);
Ok(())
}
#[test]
fn test_generator_month_start() -> Result<(), Box<dyn Error>> {
let start_date = date(2023, 10, 15); // Mid-Oct
let freq = DateFreq::MonthStart;
let n_periods = 4; // Nov, Dec, Jan, Feb
let mut generator = DatesGenerator::new(start_date, freq, n_periods)?;
assert_eq!(generator.next(), Some(date(2023, 11, 1)));
assert_eq!(generator.next(), Some(date(2023, 12, 1)));
assert_eq!(generator.next(), Some(date(2024, 1, 1)));
assert_eq!(generator.next(), Some(date(2024, 2, 1)));
assert_eq!(generator.next(), None);
Ok(())
}
#[test]
fn test_generator_month_end_leap() -> Result<(), Box<dyn Error>> {
let start_date = date(2024, 1, 31); // Jan 31
let freq = DateFreq::MonthEnd;
let n_periods = 3; // Jan, Feb (leap), Mar
let mut generator = DatesGenerator::new(start_date, freq, n_periods)?;
// find_first for Jan 31 returns Jan 31
assert_eq!(generator.next(), Some(date(2024, 1, 31)));
// find_next finds Feb 29
assert_eq!(generator.next(), Some(date(2024, 2, 29)));
// find_next finds Mar 31
assert_eq!(generator.next(), Some(date(2024, 3, 31)));
assert_eq!(generator.next(), None);
Ok(())
}
#[test]
fn test_generator_quarter_start() -> Result<(), Box<dyn Error>> {
let start_date = date(2023, 8, 1); // Mid-Q3
let freq = DateFreq::QuarterStart;
let n_periods = 3; // Q4'23, Q1'24, Q2'24
let mut generator = DatesGenerator::new(start_date, freq, n_periods)?;
assert_eq!(generator.next(), Some(date(2023, 10, 1)));
assert_eq!(generator.next(), Some(date(2024, 1, 1)));
assert_eq!(generator.next(), Some(date(2024, 4, 1)));
assert_eq!(generator.next(), None);
Ok(())
}
#[test]
fn test_generator_quarter_end() -> Result<(), Box<dyn Error>> {
let start_date = date(2023, 11, 1); // Mid-Q4
let freq = DateFreq::QuarterEnd;
let n_periods = 3; // Q4'23, Q1'24, Q2'24
let mut generator = DatesGenerator::new(start_date, freq, n_periods)?;
// find_first for Nov 1 (Q4) returns Dec 31 (Q4 end)
assert_eq!(generator.next(), Some(date(2023, 12, 31)));
// find_next finds Mar 31 (Q1 end)
assert_eq!(generator.next(), Some(date(2024, 3, 31)));
// find_next finds Jun 30 (Q2 end)
assert_eq!(generator.next(), Some(date(2024, 6, 30)));
assert_eq!(generator.next(), None);
Ok(())
}
#[test]
fn test_generator_year_start() -> Result<(), Box<dyn Error>> {
let start_date = date(2023, 1, 1); // Jan 1
let freq = DateFreq::YearStart;
let n_periods = 3; // 2023, 2024, 2025
let mut generator = DatesGenerator::new(start_date, freq, n_periods)?;
assert_eq!(generator.next(), Some(date(2023, 1, 1)));
assert_eq!(generator.next(), Some(date(2024, 1, 1)));
assert_eq!(generator.next(), Some(date(2025, 1, 1)));
assert_eq!(generator.next(), None);
Ok(())
}
#[test]
fn test_generator_year_end() -> Result<(), Box<dyn Error>> {
let start_date = date(2022, 12, 31); // Dec 31
let freq = DateFreq::YearEnd;
let n_periods = 3; // 2022, 2023, 2024
let mut generator = DatesGenerator::new(start_date, freq, n_periods)?;
// find_first for Dec 31 '22 returns Dec 31 '22
assert_eq!(generator.next(), Some(date(2022, 12, 31)));
// find_next finds Dec 31 '23
assert_eq!(generator.next(), Some(date(2023, 12, 31)));
// find_next finds Dec 31 '24
assert_eq!(generator.next(), Some(date(2024, 12, 31)));
assert_eq!(generator.next(), None);
Ok(())
}
#[test]
fn test_generator_stops_after_error_finding_next() -> Result<(), Box<dyn Error>> {
let start_year = NaiveDate::MAX.year();
let start_date = last_day_of_year(start_year - 1)?; // Dec 31 of year before MAX
let freq = DateFreq::YearEnd;
let n_periods = 3; // Try for YE(MAX-1), YE(MAX), YE(MAX+1) - last should fail
let mut generator = DatesGenerator::new(start_date, freq, n_periods)?;
// find_first returns start_date (YE MAX-1)
assert_eq!(generator.next(), Some(start_date));
// find_next finds YE(MAX)
assert_eq!(generator.next(), Some(last_day_of_year(start_year)?));
// Should be MAX
// find_next tries YE(MAX+1) - this call to find_next_date fails internally
assert_eq!(generator.next(), None);
// Returns None because internal find_next_date failed
// State after the *first* None is returned:
// Corrected assertion
assert_eq!(generator.periods_remaining, 0);
assert!(generator.next_date_candidate.is_none());
// Calling next() again should also return None
assert_eq!(generator.next(), None);
assert_eq!(generator.periods_remaining, 0);
Ok(())
}
} // end mod tests
Reference in New Issue View Git Blame Copy Permalink