From 5dc683ffaec7bfa9cc5e0e3033436120f8bc103d Mon Sep 17 00:00:00 2001 From: Palash Tyagi <23239946+Magnus167@users.noreply.github.com> Date: Sun, 20 Apr 2025 05:22:59 +0100 Subject: [PATCH 1/4] added modifications to use generator in bdates util --- src/utils/bdates.rs | 1075 ++++++++++++++++++++++++++++++++++++++----- 1 file changed, 961 insertions(+), 114 deletions(-) diff --git a/src/utils/bdates.rs b/src/utils/bdates.rs index da31d3e..7895be6 100644 --- a/src/utils/bdates.rs +++ b/src/utils/bdates.rs @@ -3,6 +3,7 @@ use std::collections::HashMap; use std::error::Error; use std::hash::Hash; use std::result::Result; +use std::str::FromStr; // Import FromStr trait /// Represents the frequency at which business dates should be generated. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] @@ -39,46 +40,8 @@ impl BDateFreq { /// /// Returns an error if the string does not match any known frequency. pub fn from_string(freq: String) -> Result> { - Self::from_str(&freq) - } - - /// Attempts to parse a frequency string slice into a `BDateFreq` enum. - /// - /// Supports various frequency codes and common aliases. - /// - /// | Code | Alias | Description | - /// |------|---------|---------------------| - /// | D | | Daily | - /// | W | WS | Weekly Monday | - /// | M | MS | Month Start | - /// | Q | QS | Quarter Start | - /// | Y | A, AS, YS | Year Start | - /// | ME | | Month End | - /// | QE | | Quarter End | - /// | WF | | Weekly Friday | - /// | YE | AE | Year End (Annual) | - /// - /// # Arguments - /// - /// * `freq` - The frequency string slice (e.g., "D", "W", "ME"). - /// - /// # Errors - /// - /// Returns an error if the string does not match any known frequency. - pub fn from_str(freq: &str) -> Result> { - let r = match freq { - "D" => BDateFreq::Daily, - "W" | "WS" => BDateFreq::WeeklyMonday, - "M" | "MS" => BDateFreq::MonthStart, - "Q" | "QS" => BDateFreq::QuarterStart, - "Y" | "A" | "AS" | "YS" => BDateFreq::YearStart, // Added Y, YS, A, AS aliases - "ME" => BDateFreq::MonthEnd, - "QE" => BDateFreq::QuarterEnd, - "WF" => BDateFreq::WeeklyFriday, - "YE" | "AE" => BDateFreq::YearEnd, // Added AE alias - _ => return Err(format!("Invalid frequency specified: {}", freq).into()), - }; - Ok(r) + // Use the FromStr implementation directly + freq.parse() } /// Returns the canonical string representation of the frequency. @@ -116,6 +79,50 @@ impl BDateFreq { } } +// Implement FromStr for BDateFreq to allow parsing directly using `parse()` +impl FromStr for BDateFreq { + type Err = Box; + + /// Attempts to parse a frequency string slice into a `BDateFreq` enum. + /// + /// Supports various frequency codes and common aliases. + /// + /// | Code | Alias | Description | + /// |------|---------|---------------------| + /// | D | | Daily | + /// | W | WS | Weekly Monday | + /// | M | MS | Month Start | + /// | Q | QS | Quarter Start | + /// | Y | A, AS, YS | Year Start | + /// | ME | | Month End | + /// | QE | | Quarter End | + /// | WF | | Weekly Friday | + /// | YE | AE | Year End (Annual) | + /// + /// # 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 { + let r = match freq { + "D" => BDateFreq::Daily, + "W" | "WS" => BDateFreq::WeeklyMonday, + "M" | "MS" => BDateFreq::MonthStart, + "Q" | "QS" => BDateFreq::QuarterStart, + "Y" | "A" | "AS" | "YS" => BDateFreq::YearStart, // Added Y, YS, A, AS aliases + "ME" => BDateFreq::MonthEnd, + "QE" => BDateFreq::QuarterEnd, + "WF" => BDateFreq::WeeklyFriday, + "YE" | "AE" => BDateFreq::YearEnd, // Added AE alias + _ => return Err(format!("Invalid frequency specified: {}", freq).into()), + }; + Ok(r) + } +} + /// Represents a list of business 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. @@ -124,6 +131,9 @@ pub struct BDatesList { start_date_str: String, end_date_str: String, freq: BDateFreq, + // Optional: Cache the generated list to avoid re-computation? + // For now, we recompute each time list(), count(), or groups() is called. + // cached_list: Option>, } // Helper enum to represent the key for grouping dates into periods. @@ -139,7 +149,7 @@ enum GroupKey { } impl BDatesList { - /// Creates a new `BDatesList` instance. + /// Creates a new `BDatesList` instance defined by a start and end date. /// /// # Arguments /// @@ -154,6 +164,52 @@ impl BDatesList { } } + /// Creates a new `BDatesList` instance defined by a start date, frequency, + /// and the number of periods (dates) to generate. + /// + /// This calculates the required dates using a `BDatesGenerator` 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 business 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: BDateFreq, + n_periods: usize, + ) -> Result> { + 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 = BDatesGenerator::new(start_date, freq, n_periods)?; + let dates: Vec = 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(BDatesList { + start_date_str, // Keep the original start date string + end_date_str, + freq, + }) + } + /// Returns the flat list of business dates within the specified range and frequency. /// /// The list is guaranteed to be sorted chronologically. @@ -250,8 +306,7 @@ impl BDatesList { /// /// # Errors /// - /// Returns a `chrono::ParseError` if the start date string is not in - /// "YYYY-MM-DD" format. + /// Returns an error if the start date string is not in "YYYY-MM-DD" format. pub fn start_date(&self) -> Result> { NaiveDate::parse_from_str(&self.start_date_str, "%Y-%m-%d").map_err(|e| e.into()) } @@ -265,8 +320,7 @@ impl BDatesList { /// /// # Errors /// - /// Returns a `chrono::ParseError` if the end date string is not in - /// "YYYY-MM-DD" format. + /// Returns an error if the end date string is not in "YYYY-MM-DD" format. pub fn end_date(&self) -> Result> { NaiveDate::parse_from_str(&self.end_date_str, "%Y-%m-%d").map_err(|e| e.into()) } @@ -287,6 +341,83 @@ impl BDatesList { } } +// --- Business Date Generator (Iterator) --- + +/// An iterator that generates a sequence of business 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. +#[derive(Debug, Clone)] +pub struct BDatesGenerator { + freq: BDateFreq, + periods_remaining: usize, + // Stores the *next* date to be yielded by the iterator. + // This is None initially or when the iterator is exhausted. + next_date_candidate: Option, +} + +impl BDatesGenerator { + /// Creates a new `BDatesGenerator`. + /// + /// It calculates the first valid business 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 business date. + /// * `freq` - The frequency for generating dates. + /// * `n_periods` - The total number of business dates to generate. + /// + /// # Errors + /// + /// Can potentially return an error if date calculations lead to overflows, + /// though this is highly unlikely with realistic date ranges. (Currently returns Ok). + /// Note: The internal `find_first_bdate_on_or_after` might panic on extreme date overflows, + /// but practical usage should be safe. + pub fn new( + start_date: NaiveDate, + freq: BDateFreq, + n_periods: usize, + ) -> Result> { + let first_date = if n_periods > 0 { + Some(find_first_bdate_on_or_after(start_date, freq)) + } else { + None // No dates to generate if n_periods is 0 + }; + + Ok(BDatesGenerator { + freq, + periods_remaining: n_periods, + next_date_candidate: first_date, + }) + } +} + +impl Iterator for BDatesGenerator { + type Item = NaiveDate; + + /// Returns the next business date in the sequence, or `None` if `n_periods` + /// dates have already been generated. + fn next(&mut self) -> Option { + // Check if exhausted or if there was no initial date + if self.periods_remaining == 0 || self.next_date_candidate.is_none() { + return None; + } + + // Get the date to return (unwrap is safe due to the check above) + let current_date = self.next_date_candidate.unwrap(); + + // Prepare the *next* candidate for the subsequent call + self.next_date_candidate = Some(find_next_bdate(current_date, self.freq)); + + // Decrement the count + self.periods_remaining -= 1; + + // Return the stored current date + Some(current_date) + } +} + // --- Internal helper functions (not part of the public API) --- /// Generates the flat list of business dates for the given range and frequency. @@ -341,6 +472,8 @@ fn get_bdates_list_with_freq( // Filter out any weekend days. While the core logic aims for business days, // this ensures robustness against edge cases where computed dates might fall // on a weekend (e.g., first day of month being Saturday). + // Note: This retain is redundant if collect_* functions are correct, but adds safety. + // It's essential for Daily, less so for others if they always return bdays. dates.retain(|d| is_weekday(*d)); // Ensure the final list is sorted. The `collect_*` functions generally @@ -351,6 +484,7 @@ fn get_bdates_list_with_freq( } /* ---------------------- Low-Level Date Collection Functions (Internal) ---------------------- */ +// These functions generate dates within a *range* [start_date, end_date] /// Returns all business days (Mon-Fri) day-by-day within the range. fn collect_daily(start_date: NaiveDate, end_date: NaiveDate) -> Vec { @@ -360,8 +494,10 @@ fn collect_daily(start_date: NaiveDate, end_date: NaiveDate) -> Vec { if is_weekday(current) { result.push(current); } - // Use succ_opt() and unwrap(), assuming valid date range and no overflow - current = current.succ_opt().unwrap(); + // Use succ_opt() and expect(), assuming valid date range and no overflow in practical scenarios + current = current + .succ_opt() + .expect("date overflow near end of supported range"); } result } @@ -379,11 +515,14 @@ fn collect_weekly( // Step through the range in 7-day increments. while current <= end_date { - result.push(current); + // Ensure the found date is actually a weekday (should be Mon/Fri but belt-and-suspenders) + if is_weekday(current) { + result.push(current); + } // Use checked_add_signed for safety, though basic addition is likely fine for 7 days. current = current .checked_add_signed(Duration::days(7)) - .expect("date overflow"); + .expect("date overflow adding 7 days"); } result } @@ -406,7 +545,7 @@ fn collect_monthly( // Iterate month by month from the start date's month up to or past the end date's month. loop { // Compute the candidate date (first or last business day) for the current month. - // Use _opt and unwrap(), expecting valid month/year combinations within realistic ranges. + // Use _opt and expect(), expecting valid month/year combinations within realistic ranges. let candidate = if want_first_day { first_business_day_of_month(year, month) } else { @@ -420,23 +559,30 @@ fn collect_monthly( // If the candidate is within the specified range [start_date, end_date], add it. if candidate >= start_date { - result.push(candidate); + // Ensure it's actually a weekday (should be, but adds safety) + if is_weekday(candidate) { + result.push(candidate); + } } // Note: We don't break if candidate < start_date because a later month's candidate // might be within the range. + // Check if the current month is the last month we should process + if year > end_date.year() || (year == end_date.year() && month >= end_date.month()) { + // If we just processed the end_date's month, stop. + // Need >= because we need to include the end date's month itself if its candidate is valid. + break; + } + // Advance to the next month. let (ny, nm) = next_month((year, month)); year = ny; month = nm; - // Optimization: Stop if we have moved clearly past the end date's year. - // If the year matches, we need to check the month. - if year > end_date.year() { - break; - } - if year == end_date.year() && month > end_date.month() { - break; + // Safety break: Stop if we have moved clearly past the end date's year. + // This check is technically redundant given the loop condition above, but harmless. + if year > end_date.year() + 1 { + break; // Avoid potential infinite loops in unexpected scenarios } } @@ -458,7 +604,7 @@ fn collect_quarterly( // Iterate quarter by quarter until we pass the end date. loop { // Compute the candidate date (first or last business day) for the current quarter. - // Use _opt and unwrap(), expecting valid quarter/year combinations. + // Use _opt and expect(), expecting valid quarter/year combinations. let candidate = if want_first_day { first_business_day_of_quarter(year, q) } else { @@ -472,11 +618,20 @@ fn collect_quarterly( // If the candidate is within the specified range [start_date, end_date], add it. if candidate >= start_date { - result.push(candidate); + // Ensure it's actually a weekday (should be, but adds safety) + if is_weekday(candidate) { + result.push(candidate); + } } // Note: We don't break if candidate < start_date because a later quarter // might be within the range. + // Check if the current quarter is the last one we should process + let end_q = month_to_quarter(end_date.month()); + if year > end_date.year() || (year == end_date.year() && q >= end_q) { + break; // Stop after processing the end_date's quarter + } + // Advance to the next quarter. if q == 4 { year += 1; @@ -484,6 +639,11 @@ fn collect_quarterly( } else { q += 1; } + + // Safety break + if year > end_date.year() + 1 { + break; + } } result @@ -502,7 +662,7 @@ fn collect_yearly( // Iterate year by year until we pass the end date's year. while year <= end_date.year() { // Compute the candidate date (first or last business day) for the current year. - // Use _opt and unwrap(), expecting valid year. + // Use _opt and expect(), expecting valid year. let candidate = if want_first_day { first_business_day_of_year(year) } else { @@ -511,10 +671,14 @@ fn collect_yearly( // If the candidate is within the specified range [start_date, end_date], add it. if candidate >= start_date && candidate <= end_date { - result.push(candidate); - } else if candidate > end_date { - // Optimization: If the candidate for the current year is already past end_date, - // then candidates for all subsequent years will also be past end_date. + // Ensure it's actually a weekday (should be, but adds safety) + if is_weekday(candidate) { + result.push(candidate); + } + } else if want_first_day && candidate > end_date { + // Optimization: If the *first* bday of the year is already past end_date, + // no subsequent year's first bday will be in range. + // Similar logic applies for last bday if candidate > end_date, but it's less likely to trigger early. break; } // Note: We don't break if candidate < start_date because a later year's candidate @@ -538,20 +702,22 @@ fn is_weekday(date: NaiveDate) -> bool { fn move_to_weekday_on_or_after(date: NaiveDate, target: Weekday) -> NaiveDate { let mut current = date; while current.weekday() != target { - // Use succ_opt() and unwrap(), assuming valid date and no overflow - current = current.succ_opt().unwrap(); + // Use succ_opt() and expect(), assuming valid date and no overflow + current = current + .succ_opt() + .expect("date overflow moving to next weekday"); } current } /// Return the earliest business day of the given (year, month). fn first_business_day_of_month(year: i32, month: u32) -> NaiveDate { - // Start with the 1st of the month. Use _opt and unwrap(), assuming valid Y/M. + // Start with the 1st of the month. Use _opt and expect(), assuming valid Y/M. let mut d = NaiveDate::from_ymd_opt(year, month, 1).expect("invalid year-month combination"); // If it’s Sat/Sun, move forward until we find a weekday. while !is_weekday(d) { - // Use succ_opt() and unwrap(), assuming valid date and no overflow. - d = d.succ_opt().unwrap(); + // Use succ_opt() and expect(), assuming valid date and no overflow. + d = d.succ_opt().expect("date overflow finding first bday"); } d } @@ -559,13 +725,13 @@ fn first_business_day_of_month(year: i32, month: u32) -> NaiveDate { /// Return the latest business day of the given (year, month). fn last_business_day_of_month(year: i32, month: u32) -> NaiveDate { let last_dom = days_in_month(year, month); - // Use _opt and unwrap(), assuming valid Y/M/D combination. + // Use _opt and expect(), assuming valid Y/M/D combination. let mut d = NaiveDate::from_ymd_opt(year, month, last_dom).expect("invalid year-month-day combination"); // If it’s Sat/Sun, move backward until we find a weekday. while !is_weekday(d) { - // Use pred_opt() and unwrap(), assuming valid date and no underflow. - d = d.pred_opt().unwrap(); + // Use pred_opt() and expect(), assuming valid date and no underflow. + d = d.pred_opt().expect("date underflow finding last bday"); } d } @@ -579,9 +745,9 @@ fn days_in_month(year: i32, month: u32) -> u32 { } else { (year, month + 1) }; - // Use _opt and unwrap(), assuming valid Y/M combination (start of next month). + // Use _opt and expect(), assuming valid Y/M combination (start of next month). let first_of_next = NaiveDate::from_ymd_opt(ny, nm, 1).expect("invalid next year-month"); - // Use pred_opt() and unwrap(), assuming valid date and no underflow (first of month - 1). + // Use pred_opt() and expect(), assuming valid date and no underflow (first of month - 1). let last_of_this = first_of_next .pred_opt() .expect("invalid date before first of month"); @@ -590,65 +756,231 @@ fn days_in_month(year: i32, month: u32) -> u32 { /// Converts a month number (1-12) to a quarter number (1-4). fn month_to_quarter(m: u32) -> u32 { - (m - 1) / 3 + 1 // Simple integer division for mapping + match m { + 1..=3 => 1, + 4..=6 => 2, + 7..=9 => 3, + 10..=12 => 4, + _ => panic!("Invalid month: {}", m), // Should not happen with valid dates + } } /// Returns the 1st day of the month that starts a given (year, quarter). -fn quarter_to_first_date(year: i32, quarter: u32) -> NaiveDate { - let month = match quarter { +fn quarter_start_month(quarter: u32) -> u32 { + match quarter { 1 => 1, 2 => 4, 3 => 7, 4 => 10, _ => panic!("invalid quarter: {}", quarter), // This function expects quarter 1-4 - }; - // Use _opt and unwrap(), assuming valid Y/M/D combination (first day of quarter month). - NaiveDate::from_ymd_opt(year, month, 1).expect("invalid year/month derived from quarter") + } } /// Return the earliest business day in the given (year, quarter). fn first_business_day_of_quarter(year: i32, quarter: u32) -> NaiveDate { - let mut d = quarter_to_first_date(year, quarter); - // If the first day is a weekend, move forward to the next weekday. - while !is_weekday(d) { - // Use succ_opt() and unwrap(), assuming valid date and no overflow. - d = d.succ_opt().unwrap(); - } - d + let month = quarter_start_month(quarter); + first_business_day_of_month(year, month) } /// Return the last business day in the given (year, quarter). fn last_business_day_of_quarter(year: i32, quarter: u32) -> NaiveDate { // The last month of a quarter is the start month + 2. - let start = quarter_to_first_date(year, quarter); - let last_month_in_quarter = start.month() + 2; + let last_month_in_quarter = match quarter { + 1 => 3, + 2 => 6, + 3 => 9, + 4 => 12, + _ => panic!("invalid quarter: {}", quarter), + }; last_business_day_of_month(year, last_month_in_quarter) } /// Returns the earliest business day (Mon-Fri) of the given year. fn first_business_day_of_year(year: i32) -> NaiveDate { - // Start with Jan 1st. Use _opt and unwrap(), assuming valid Y/M/D combination. + // Start with Jan 1st. Use _opt and expect(), assuming valid Y/M/D combination. let mut d = NaiveDate::from_ymd_opt(year, 1, 1).expect("invalid year for Jan 1st"); // If Jan 1st is a weekend, move forward to the next weekday. while !is_weekday(d) { - // Use succ_opt() and unwrap(), assuming valid date and no overflow. - d = d.succ_opt().unwrap(); + // Use succ_opt() and expect(), assuming valid date and no overflow. + d = d + .succ_opt() + .expect("date overflow finding first bday of year"); } d } /// Returns the last business day (Mon-Fri) of the given year. fn last_business_day_of_year(year: i32) -> NaiveDate { - // Start with Dec 31st. Use _opt and unwrap(), assuming valid Y/M/D combination. + // Start with Dec 31st. Use _opt and expect(), assuming valid Y/M/D combination. let mut d = NaiveDate::from_ymd_opt(year, 12, 31).expect("invalid year for Dec 31st"); // If Dec 31st is a weekend, move backward to the previous weekday. while !is_weekday(d) { - // Use pred_opt() and unwrap(), assuming valid date and no underflow. - d = d.pred_opt().unwrap(); + // Use pred_opt() and expect(), assuming valid date and no underflow. + d = d + .pred_opt() + .expect("date underflow finding last bday of year"); } d } +// --- Generator Helper Functions --- + +/// Finds the *first* valid business date according to the frequency, +/// starting the search *on or after* the given `start_date`. +/// Panics on date overflow/underflow in extreme cases, but generally safe. +fn find_first_bdate_on_or_after(start_date: NaiveDate, freq: BDateFreq) -> NaiveDate { + match freq { + BDateFreq::Daily => { + let mut d = start_date; + while !is_weekday(d) { + d = d + .succ_opt() + .expect("Date overflow finding first daily date"); + } + d + } + BDateFreq::WeeklyMonday => move_to_weekday_on_or_after(start_date, Weekday::Mon), + BDateFreq::WeeklyFriday => move_to_weekday_on_or_after(start_date, Weekday::Fri), + BDateFreq::MonthStart => { + let mut candidate = first_business_day_of_month(start_date.year(), start_date.month()); + if candidate < start_date { + // If the first bday of the current month is before start_date, + // we need the first bday of the *next* month. + let (next_y, next_m) = if start_date.month() == 12 { + (start_date.year() + 1, 1) + } else { + (start_date.year(), start_date.month() + 1) + }; + candidate = first_business_day_of_month(next_y, next_m); + } + candidate + } + BDateFreq::MonthEnd => { + let mut candidate = last_business_day_of_month(start_date.year(), start_date.month()); + if candidate < start_date { + // If the last bday of current month is before start_date, + // we need the last bday of the *next* month. + let (next_y, next_m) = if start_date.month() == 12 { + (start_date.year() + 1, 1) + } else { + (start_date.year(), start_date.month() + 1) + }; + candidate = last_business_day_of_month(next_y, next_m); + } + candidate + } + BDateFreq::QuarterStart => { + let current_q = month_to_quarter(start_date.month()); + let mut candidate = first_business_day_of_quarter(start_date.year(), current_q); + if candidate < start_date { + // If the first bday of the current quarter is before start_date, + // we need the first bday of the *next* quarter. + let (next_y, next_q) = if current_q == 4 { + (start_date.year() + 1, 1) + } else { + (start_date.year(), current_q + 1) + }; + candidate = first_business_day_of_quarter(next_y, next_q); + } + candidate + } + BDateFreq::QuarterEnd => { + let current_q = month_to_quarter(start_date.month()); + let mut candidate = last_business_day_of_quarter(start_date.year(), current_q); + if candidate < start_date { + // If the last bday of the current quarter is before start_date, + // we need the last bday of the *next* quarter. + let (next_y, next_q) = if current_q == 4 { + (start_date.year() + 1, 1) + } else { + (start_date.year(), current_q + 1) + }; + candidate = last_business_day_of_quarter(next_y, next_q); + } + candidate + } + BDateFreq::YearStart => { + let mut candidate = first_business_day_of_year(start_date.year()); + if candidate < start_date { + // If the first bday of the current year is before start_date, + // we need the first bday of the *next* year. + candidate = first_business_day_of_year(start_date.year() + 1); + } + candidate + } + BDateFreq::YearEnd => { + let mut candidate = last_business_day_of_year(start_date.year()); + if candidate < start_date { + // If the last bday of the current year is before start_date, + // we need the last bday of the *next* year. + candidate = last_business_day_of_year(start_date.year() + 1); + } + candidate + } + } +} + +/// Finds the *next* valid business date according to the frequency, +/// given the `current_date` (which is assumed to be a valid date previously generated). +/// Panics on date overflow/underflow in extreme cases, but generally safe. +fn find_next_bdate(current_date: NaiveDate, freq: BDateFreq) -> NaiveDate { + match freq { + BDateFreq::Daily => { + let mut next_day = current_date + .succ_opt() + .expect("Date overflow finding next daily"); + while !is_weekday(next_day) { + next_day = next_day + .succ_opt() + .expect("Date overflow finding next daily weekday"); + } + next_day + } + BDateFreq::WeeklyMonday | BDateFreq::WeeklyFriday => { + // Assuming current_date is already a Mon/Fri, the next one is 7 days later. + current_date + .checked_add_signed(Duration::days(7)) + .expect("Date overflow adding 7 days") + } + BDateFreq::MonthStart => { + let (next_y, next_m) = if current_date.month() == 12 { + (current_date.year() + 1, 1) + } else { + (current_date.year(), current_date.month() + 1) + }; + first_business_day_of_month(next_y, next_m) + } + BDateFreq::MonthEnd => { + let (next_y, next_m) = if current_date.month() == 12 { + (current_date.year() + 1, 1) + } else { + (current_date.year(), current_date.month() + 1) + }; + last_business_day_of_month(next_y, next_m) + } + BDateFreq::QuarterStart => { + let current_q = month_to_quarter(current_date.month()); + let (next_y, next_q) = if current_q == 4 { + (current_date.year() + 1, 1) + } else { + (current_date.year(), current_q + 1) + }; + first_business_day_of_quarter(next_y, next_q) + } + BDateFreq::QuarterEnd => { + let current_q = month_to_quarter(current_date.month()); + let (next_y, next_q) = if current_q == 4 { + (current_date.year() + 1, 1) + } else { + (current_date.year(), current_q + 1) + }; + last_business_day_of_quarter(next_y, next_q) + } + BDateFreq::YearStart => first_business_day_of_year(current_date.year() + 1), + BDateFreq::YearEnd => last_business_day_of_year(current_date.year() + 1), + } +} + // --- Example Usage and Tests --- #[cfg(test)] @@ -656,9 +988,9 @@ mod tests { use super::*; use chrono::NaiveDate; - // Helper to create a NaiveDate for tests, handling the unwrap for fixed dates. + // Helper to create a NaiveDate for tests, handling the expect for fixed dates. fn date(year: i32, month: u32, day: u32) -> NaiveDate { - NaiveDate::from_ymd_opt(year, month, day).unwrap() + NaiveDate::from_ymd_opt(year, month, day).expect("Invalid date in test setup") } // --- BDateFreq Tests --- @@ -666,6 +998,7 @@ mod tests { #[test] fn test_bdatefreq_from_str() -> Result<(), Box> { assert_eq!(BDateFreq::from_str("D")?, BDateFreq::Daily); + assert_eq!("D".parse::()?, BDateFreq::Daily); // Test FromStr impl assert_eq!(BDateFreq::from_str("W")?, BDateFreq::WeeklyMonday); assert_eq!(BDateFreq::from_str("M")?, BDateFreq::MonthStart); assert_eq!(BDateFreq::from_str("Q")?, BDateFreq::QuarterStart); @@ -675,10 +1008,12 @@ mod tests { assert_eq!(BDateFreq::from_str("A")?, BDateFreq::YearStart); assert_eq!(BDateFreq::from_str("AS")?, BDateFreq::YearStart); assert_eq!(BDateFreq::from_str("YS")?, BDateFreq::YearStart); + assert_eq!("Y".parse::()?, BDateFreq::YearStart); // Test FromStr impl assert_eq!(BDateFreq::from_str("ME")?, BDateFreq::MonthEnd); assert_eq!(BDateFreq::from_str("QE")?, BDateFreq::QuarterEnd); assert_eq!(BDateFreq::from_str("WF")?, BDateFreq::WeeklyFriday); + assert_eq!("WF".parse::()?, BDateFreq::WeeklyFriday); // Test FromStr impl // Test YearEnd codes and aliases (YE, AE) assert_eq!(BDateFreq::from_str("YE")?, BDateFreq::YearEnd); @@ -691,6 +1026,7 @@ mod tests { // Test invalid string assert!(BDateFreq::from_str("INVALID").is_err()); + assert!("INVALID".parse::().is_err()); // Test FromStr impl let err = BDateFreq::from_str("INVALID").unwrap_err(); assert_eq!(err.to_string(), "Invalid frequency specified: INVALID"); @@ -734,7 +1070,7 @@ mod tests { // --- BDatesList Property Tests --- #[test] - fn test_bdates_list_properties() -> Result<(), Box> { + fn test_bdates_list_properties_new() -> Result<(), Box> { let start_str = "2023-01-01".to_string(); let end_str = "2023-12-31".to_string(); let freq = BDateFreq::QuarterEnd; @@ -757,7 +1093,72 @@ mod tests { } #[test] - fn test_bdates_list_invalid_date_string() { + fn test_bdates_list_properties_from_n_periods() -> Result<(), Box> { + let start_str = "2023-01-01".to_string(); // Sunday + let freq = BDateFreq::Daily; + let n_periods = 5; // Expect: Jan 2, 3, 4, 5, 6 + let dates_list = BDatesList::from_n_periods(start_str.clone(), freq, n_periods)?; + + // check start_date_str (should be original) + assert_eq!(dates_list.start_date_str(), start_str); + // check end_date_str (should be the last generated date) + assert_eq!(dates_list.end_date_str(), "2023-01-06"); + // check frequency enum + assert_eq!(dates_list.freq(), freq); + // check frequency string + assert_eq!(dates_list.freq_str(), "D"); + + // Check parsed dates + assert_eq!(dates_list.start_date()?, date(2023, 1, 1)); + assert_eq!(dates_list.end_date()?, date(2023, 1, 6)); + + // Check the actual list matches + assert_eq!( + dates_list.list()?, + vec![ + date(2023, 1, 2), + date(2023, 1, 3), + date(2023, 1, 4), + date(2023, 1, 5), + date(2023, 1, 6) + ] + ); + assert_eq!(dates_list.count()?, 5); + + Ok(()) + } + + #[test] + fn test_bdates_list_from_n_periods_zero_periods() { + let start_str = "2023-01-01".to_string(); + let freq = BDateFreq::Daily; + let n_periods = 0; + let result = BDatesList::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] + fn test_bdates_list_from_n_periods_invalid_start_date() { + let start_str = "invalid-date".to_string(); + let freq = BDateFreq::Daily; + let n_periods = 5; + let result = BDatesList::from_n_periods(start_str.clone(), freq, n_periods); + assert!(result.is_err()); + // Error comes from NaiveDate::parse_from_str + assert!( + result + .unwrap_err() + .to_string() + .contains("input contains invalid characters") + ); + } + + #[test] + fn test_bdates_list_invalid_date_string_new() { let dates_list_start_invalid = BDatesList::new( "invalid-date".to_string(), "2023-12-31".to_string(), @@ -1034,6 +1435,7 @@ mod tests { // 2023-10-02 (Q4 2023 start - Oct 1st is Sunday) -> Included // 2024-01-01 (Q1 2024 start - Jan 1st is Monday) -> Included // 2024-04-01 (Q2 2024 start) -> Included + // 2024-07-01 (Q3 2024 start) -> Excluded by end_date 2024-05-01 // Expected groups: Q4 2023, Q1 2024, Q2 2024 assert_eq!(groups.len(), 3); @@ -1177,15 +1579,15 @@ mod tests { } #[test] - /// Tests `quarter_to_first_date` for all quarters. - fn test_quarter_to_first_date() { - assert_eq!(quarter_to_first_date(2023, 1), date(2023, 1, 1)); - assert_eq!(quarter_to_first_date(2023, 2), date(2023, 4, 1)); - assert_eq!(quarter_to_first_date(2023, 3), date(2023, 7, 1)); - assert_eq!(quarter_to_first_date(2023, 4), date(2023, 10, 1)); - // Panics on invalid quarter - let result = std::panic::catch_unwind(|| quarter_to_first_date(2023, 5)); - assert!(result.is_err()); + #[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] @@ -1360,7 +1762,7 @@ mod tests { // Q starts >= start_date AND <= end_date: Q4 2023, Q1 2024 // Q3 2023 start bday (Jul 3rd) < start_date (Aug 15th) -> Excluded // Q4 2023 start bday (Oct 2nd) >= start_date (Aug 15th) -> Included - // Q1 2024 start bday (Jan 1st) >= start_date (Aug 15th) -> Included + // Q1 2024 start bday (Jan 1st) >= start_date (Aug 15th) AND <= end_date -> Included // Q2 2024 start bday (Apr 1st) > end_date (Feb 15th) -> Excluded assert_eq!( collect_quarterly(start, end, true), @@ -1369,7 +1771,7 @@ mod tests { // Q ends >= start_date AND <= end_date: Q3 2023, Q4 2023 // Q3 2023 end bday (Sep 29th) >= start_date (Aug 15th) -> Included // Q4 2023 end bday (Dec 29th) >= start_date (Aug 15th) -> Included - // Q1 2024 end bday (Mar 31st) > end_date (Feb 15th) -> Excluded + // Q1 2024 end bday (Mar 29th) > end_date (Feb 15th) -> Excluded assert_eq!( collect_quarterly(start, end, false), vec![date(2023, 9, 29), date(2023, 12, 29)] @@ -1397,6 +1799,38 @@ mod tests { assert_eq!(collect_quarterly(start, end, false), vec![]); } + #[test] + fn test_collect_quarterly_full_year_start() { + let start = date(2023, 1, 1); + let end = date(2023, 12, 31); + // Q1: Jan 2, Q2: Apr 3, Q3: Jul 3, Q4: Oct 2 + assert_eq!( + collect_quarterly(start, end, true), + vec![ + date(2023, 1, 2), + date(2023, 4, 3), + date(2023, 7, 3), + date(2023, 10, 2) + ] + ); + } + + #[test] + fn test_collect_quarterly_full_year_end() { + let start = date(2023, 1, 1); + let end = date(2023, 12, 31); + // Q1: Mar 31, Q2: Jun 30, Q3: Sep 29, Q4: Dec 29 + assert_eq!( + collect_quarterly(start, end, false), + vec![ + date(2023, 3, 31), + date(2023, 6, 30), + date(2023, 9, 29), + date(2023, 12, 29) + ] + ); + } + // Test `collect_yearly` edge cases #[test] fn test_collect_yearly_range_starts_mid_year_ends_mid_year() -> Result<(), Box> { @@ -1404,12 +1838,12 @@ mod tests { let end = date(2024, 6, 1); // Mid 2024 // Year starts >= start_date AND <= end_date: 2024 // 2023 start bday (Jan 2nd) < start_date (Jun 1st) -> Excluded - // 2024 start bday (Jan 1st) >= start_date (Jun 1st) -> Included + // 2024 start bday (Jan 1st) >= start_date (Jun 1st) AND <= end_date -> Included // 2025 start bday (Jan 1st) > end_date (Jun 1st) -> Excluded assert_eq!(collect_yearly(start, end, true), vec![date(2024, 1, 1)]); // Year ends >= start_date AND <= end_date: 2023 // 2023 end bday (Dec 29th) >= start_date (Jun 1st) -> Included - // 2024 end bday (Dec 31st) > end_date (Jun 1st) -> Included + // 2024 end bday (Dec 31st) > end_date (Jun 1st) -> Excluded <-- Correction: Original thought was wrong assert_eq!(collect_yearly(start, end, false), vec![date(2023, 12, 29)]); Ok(()) } @@ -1447,4 +1881,417 @@ mod tests { vec![date(2022, 12, 30), date(2023, 12, 29), date(2024, 12, 31)] ); } -} + + // --- Tests for Generator Helper Functions --- + + #[test] + fn test_find_first_bdate_on_or_after() { + // Daily + assert_eq!( + find_first_bdate_on_or_after(date(2023, 11, 8), BDateFreq::Daily), + date(2023, 11, 8) + ); // Wed -> Wed + assert_eq!( + find_first_bdate_on_or_after(date(2023, 11, 11), BDateFreq::Daily), + date(2023, 11, 13) + ); // Sat -> Mon + + // Weekly Mon + assert_eq!( + find_first_bdate_on_or_after(date(2023, 11, 8), BDateFreq::WeeklyMonday), + date(2023, 11, 13) + ); // Wed -> Next Mon + assert_eq!( + find_first_bdate_on_or_after(date(2023, 11, 13), BDateFreq::WeeklyMonday), + date(2023, 11, 13) + ); // Mon -> Mon + assert_eq!( + find_first_bdate_on_or_after(date(2023, 11, 12), BDateFreq::WeeklyMonday), + date(2023, 11, 13) + ); // Sun -> Mon + + // Weekly Fri + assert_eq!( + find_first_bdate_on_or_after(date(2023, 11, 8), BDateFreq::WeeklyFriday), + date(2023, 11, 10) + ); // Wed -> Fri + assert_eq!( + find_first_bdate_on_or_after(date(2023, 11, 10), BDateFreq::WeeklyFriday), + date(2023, 11, 10) + ); // Fri -> Fri + assert_eq!( + find_first_bdate_on_or_after(date(2023, 11, 11), BDateFreq::WeeklyFriday), + date(2023, 11, 17) + ); // Sat -> Next Fri + + // Month Start + assert_eq!( + find_first_bdate_on_or_after(date(2023, 11, 1), BDateFreq::MonthStart), + date(2023, 11, 1) + ); // Nov 1 (Wed) -> Nov 1 + assert_eq!( + find_first_bdate_on_or_after(date(2023, 10, 15), BDateFreq::MonthStart), + date(2023, 11, 1) + ); // Mid Oct -> Nov 1 + assert_eq!( + find_first_bdate_on_or_after(date(2023, 12, 15), BDateFreq::MonthStart), + date(2024, 1, 1) + ); // Mid Dec -> Jan 1 + assert_eq!( + find_first_bdate_on_or_after(date(2023, 10, 1), BDateFreq::MonthStart), + date(2023, 10, 2) + ); // Oct 1 (Sun) -> Oct 2 + + // Month End + assert_eq!( + find_first_bdate_on_or_after(date(2023, 11, 30), BDateFreq::MonthEnd), + date(2023, 11, 30) + ); // Nov 30 (Thu) -> Nov 30 + assert_eq!( + find_first_bdate_on_or_after(date(2023, 11, 15), BDateFreq::MonthEnd), + date(2023, 11, 30) + ); // Mid Nov -> Nov 30 + assert_eq!( + find_first_bdate_on_or_after(date(2023, 12, 30), BDateFreq::MonthEnd), + date(2024, 1, 31) + ); // Dec 30 (Sat) -> Jan 31 (Dec end was 29th, which is < 30th) + assert_eq!( + find_first_bdate_on_or_after(date(2023, 12, 29), BDateFreq::MonthEnd), + date(2023, 12, 29) + ); // Dec 29 (Fri) -> Dec 29 + assert_eq!( + find_first_bdate_on_or_after(date(2023, 9, 30), BDateFreq::MonthEnd), + date(2023, 10, 31) + ); // Sep 30 (Sat) -> Oct 31 (Sep end was 29th, < 30th) + + // Quarter Start + assert_eq!( + find_first_bdate_on_or_after(date(2023, 10, 2), BDateFreq::QuarterStart), + date(2023, 10, 2) + ); // Q4 Start (Mon) -> Q4 Start + assert_eq!( + find_first_bdate_on_or_after(date(2023, 8, 15), BDateFreq::QuarterStart), + date(2023, 10, 2) + ); // Mid Q3 -> Q4 Start + assert_eq!( + find_first_bdate_on_or_after(date(2023, 11, 15), BDateFreq::QuarterStart), + date(2024, 1, 1) + ); // Mid Q4 -> Q1 Start + assert_eq!( + find_first_bdate_on_or_after(date(2023, 1, 1), BDateFreq::QuarterStart), + date(2023, 1, 2) + ); // Jan 1 (Sun) -> Jan 2 (Mon) + + // Quarter End + assert_eq!( + find_first_bdate_on_or_after(date(2023, 9, 29), BDateFreq::QuarterEnd), + date(2023, 9, 29) + ); // Q3 End (Fri) -> Q3 End + assert_eq!( + find_first_bdate_on_or_after(date(2023, 8, 15), BDateFreq::QuarterEnd), + date(2023, 9, 29) + ); // Mid Q3 -> Q3 End + assert_eq!( + find_first_bdate_on_or_after(date(2023, 10, 15), BDateFreq::QuarterEnd), + date(2023, 12, 29) + ); // Mid Q4 -> Q4 End + assert_eq!( + find_first_bdate_on_or_after(date(2023, 12, 30), BDateFreq::QuarterEnd), + date(2024, 3, 29) + ); // Dec 30 (Sat) -> Q1 End (Q4 end was 29th, < 30th) + + // Year Start + assert_eq!( + find_first_bdate_on_or_after(date(2024, 1, 1), BDateFreq::YearStart), + date(2024, 1, 1) + ); // Jan 1 (Mon) -> Jan 1 + assert_eq!( + find_first_bdate_on_or_after(date(2023, 6, 15), BDateFreq::YearStart), + date(2024, 1, 1) + ); // Mid 2023 -> Jan 1 2024 + assert_eq!( + find_first_bdate_on_or_after(date(2023, 1, 1), BDateFreq::YearStart), + date(2023, 1, 2) + ); // Jan 1 (Sun) -> Jan 2 + + // Year End + assert_eq!( + find_first_bdate_on_or_after(date(2023, 12, 29), BDateFreq::YearEnd), + date(2023, 12, 29) + ); // Dec 29 (Fri) -> Dec 29 + assert_eq!( + find_first_bdate_on_or_after(date(2023, 6, 15), BDateFreq::YearEnd), + date(2023, 12, 29) + ); // Mid 2023 -> Dec 29 2023 + assert_eq!( + find_first_bdate_on_or_after(date(2023, 12, 30), BDateFreq::YearEnd), + date(2024, 12, 31) + ); // Dec 30 (Sat) -> Dec 31 2024 (2023 end was 29th, < 30th) + } + + #[test] + fn test_find_next_bdate() { + // Daily + assert_eq!( + find_next_bdate(date(2023, 11, 8), BDateFreq::Daily), + date(2023, 11, 9) + ); // Wed -> Thu + assert_eq!( + find_next_bdate(date(2023, 11, 10), BDateFreq::Daily), + date(2023, 11, 13) + ); // Fri -> Mon + + // Weekly Mon + assert_eq!( + find_next_bdate(date(2023, 11, 13), BDateFreq::WeeklyMonday), + date(2023, 11, 20) + ); // Mon -> Next Mon + + // Weekly Fri + assert_eq!( + find_next_bdate(date(2023, 11, 10), BDateFreq::WeeklyFriday), + date(2023, 11, 17) + ); // Fri -> Next Fri + + // Month Start + assert_eq!( + find_next_bdate(date(2023, 11, 1), BDateFreq::MonthStart), + date(2023, 12, 1) + ); // Nov 1 -> Dec 1 + assert_eq!( + find_next_bdate(date(2023, 12, 1), BDateFreq::MonthStart), + date(2024, 1, 1) + ); // Dec 1 -> Jan 1 + + // Month End + assert_eq!( + find_next_bdate(date(2023, 10, 31), BDateFreq::MonthEnd), + date(2023, 11, 30) + ); // Oct 31 -> Nov 30 + assert_eq!( + find_next_bdate(date(2023, 11, 30), BDateFreq::MonthEnd), + date(2023, 12, 29) + ); // Nov 30 -> Dec 29 + assert_eq!( + find_next_bdate(date(2023, 12, 29), BDateFreq::MonthEnd), + date(2024, 1, 31) + ); // Dec 29 -> Jan 31 + + // Quarter Start + assert_eq!( + find_next_bdate(date(2023, 10, 2), BDateFreq::QuarterStart), + date(2024, 1, 1) + ); // Q4 Start -> Q1 Start + assert_eq!( + find_next_bdate(date(2024, 1, 1), BDateFreq::QuarterStart), + date(2024, 4, 1) + ); // Q1 Start -> Q2 Start + + // Quarter End + assert_eq!( + find_next_bdate(date(2023, 9, 29), BDateFreq::QuarterEnd), + date(2023, 12, 29) + ); // Q3 End -> Q4 End + assert_eq!( + find_next_bdate(date(2023, 12, 29), BDateFreq::QuarterEnd), + date(2024, 3, 29) + ); // Q4 End -> Q1 End (Mar 31 2024 is Sun) + + // Year Start + assert_eq!( + find_next_bdate(date(2023, 1, 2), BDateFreq::YearStart), + date(2024, 1, 1) + ); // 2023 Start -> 2024 Start + assert_eq!( + find_next_bdate(date(2024, 1, 1), BDateFreq::YearStart), + date(2025, 1, 1) + ); // 2024 Start -> 2025 Start + + // Year End + assert_eq!( + find_next_bdate(date(2022, 12, 30), BDateFreq::YearEnd), + date(2023, 12, 29) + ); // 2022 End -> 2023 End + assert_eq!( + find_next_bdate(date(2023, 12, 29), BDateFreq::YearEnd), + date(2024, 12, 31) + ); // 2023 End -> 2024 End + } + + // --- Tests for BDatesGenerator --- + + #[test] + fn test_generator_new_zero_periods() -> Result<(), Box> { + let start_date = date(2023, 1, 1); + let freq = BDateFreq::Daily; + let n_periods = 0; + let mut generator = BDatesGenerator::new(start_date, freq, n_periods)?; + assert_eq!(generator.next(), None); // Should be immediately exhausted + Ok(()) + } + + #[test] + fn test_generator_daily() -> Result<(), Box> { + let start_date = date(2023, 11, 10); // Friday + let freq = BDateFreq::Daily; + let n_periods = 4; + let mut generator = BDatesGenerator::new(start_date, freq, n_periods)?; + + assert_eq!(generator.next(), Some(date(2023, 11, 10))); // Fri + assert_eq!(generator.next(), Some(date(2023, 11, 13))); // Mon + assert_eq!(generator.next(), Some(date(2023, 11, 14))); // Tue + assert_eq!(generator.next(), Some(date(2023, 11, 15))); // Wed + assert_eq!(generator.next(), None); // Exhausted + + // Test starting on weekend + let start_date_sat = date(2023, 11, 11); // Saturday + let mut generator_sat = BDatesGenerator::new(start_date_sat, freq, 2)?; + assert_eq!(generator_sat.next(), Some(date(2023, 11, 13))); // Mon + assert_eq!(generator_sat.next(), Some(date(2023, 11, 14))); // Tue + assert_eq!(generator_sat.next(), None); + + Ok(()) + } + + #[test] + fn test_generator_weekly_monday() -> Result<(), Box> { + let start_date = date(2023, 11, 8); // Wednesday + let freq = BDateFreq::WeeklyMonday; + let n_periods = 3; + let mut generator = BDatesGenerator::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> { + let start_date = date(2023, 11, 11); // Saturday + let freq = BDateFreq::WeeklyFriday; + let n_periods = 3; + let mut generator = BDatesGenerator::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> { + let start_date = date(2023, 10, 15); // Mid-Oct + let freq = BDateFreq::MonthStart; + let n_periods = 4; // Nov, Dec, Jan, Feb + let mut generator = BDatesGenerator::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() -> Result<(), Box> { + let start_date = date(2023, 9, 30); // Sep 30 (Sat) + let freq = BDateFreq::MonthEnd; + let n_periods = 4; // Oct, Nov, Dec, Jan + let mut generator = BDatesGenerator::new(start_date, freq, n_periods)?; + + assert_eq!(generator.next(), Some(date(2023, 10, 31))); // Sep end was 29th < 30th, so start with Oct end + assert_eq!(generator.next(), Some(date(2023, 11, 30))); + assert_eq!(generator.next(), Some(date(2023, 12, 29))); + assert_eq!(generator.next(), Some(date(2024, 1, 31))); + assert_eq!(generator.next(), None); + + Ok(()) + } + + #[test] + fn test_generator_quarter_start() -> Result<(), Box> { + let start_date = date(2023, 8, 1); // Mid-Q3 + let freq = BDateFreq::QuarterStart; + let n_periods = 3; // Q4'23, Q1'24, Q2'24 + let mut generator = BDatesGenerator::new(start_date, freq, n_periods)?; + + assert_eq!(generator.next(), Some(date(2023, 10, 2))); // Q3 start was Jul 3, < Aug 1. Next is Q4 start. + 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> { + let start_date = date(2023, 11, 1); // Mid-Q4 + let freq = BDateFreq::QuarterEnd; + let n_periods = 3; // Q4'23, Q1'24, Q2'24 + let mut generator = BDatesGenerator::new(start_date, freq, n_periods)?; + + assert_eq!(generator.next(), Some(date(2023, 12, 29))); // Q4 end is Dec 29 >= Nov 1 + assert_eq!(generator.next(), Some(date(2024, 3, 29))); // Q1 end (Mar 31 is Sun) + assert_eq!(generator.next(), Some(date(2024, 6, 28))); // Q2 end (Jun 30 is Sun) + assert_eq!(generator.next(), None); + + Ok(()) + } + + #[test] + fn test_generator_year_start() -> Result<(), Box> { + let start_date = date(2023, 1, 1); // Jan 1 (Sun) + let freq = BDateFreq::YearStart; + let n_periods = 3; // 2023, 2024, 2025 + let mut generator = BDatesGenerator::new(start_date, freq, n_periods)?; + + assert_eq!(generator.next(), Some(date(2023, 1, 2))); // 2023 start bday >= Jan 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> { + let start_date = date(2022, 12, 31); // Dec 31 (Sat) + let freq = BDateFreq::YearEnd; + let n_periods = 3; // 2023, 2024, 2025 + let mut generator = BDatesGenerator::new(start_date, freq, n_periods)?; + + assert_eq!(generator.next(), Some(date(2023, 12, 29))); // 2022 end was Dec 30 < Dec 31. Next is 2023 end. + assert_eq!(generator.next(), Some(date(2024, 12, 31))); + assert_eq!(generator.next(), Some(date(2025, 12, 31))); + assert_eq!(generator.next(), None); + + Ok(()) + } + + #[test] + fn test_generator_collect() -> Result<(), Box> { + let start_date = date(2023, 11, 10); // Friday + let freq = BDateFreq::Daily; + let n_periods = 4; + let generator = BDatesGenerator::new(start_date, freq, n_periods)?; // Use non-mut binding for collect + let dates: Vec = generator.collect(); + + assert_eq!( + dates, + vec![ + date(2023, 11, 10), // Fri + date(2023, 11, 13), // Mon + date(2023, 11, 14), // Tue + date(2023, 11, 15) // Wed + ] + ); + Ok(()) + } +} // end mod tests From 40f5f68e65d73548ffa29170e432727a9ee30d1d Mon Sep 17 00:00:00 2001 From: Palash Tyagi <23239946+Magnus167@users.noreply.github.com> Date: Sun, 20 Apr 2025 05:23:30 +0100 Subject: [PATCH 2/4] added docstring for bdateslist --- src/utils/bdates.rs | 94 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 94 insertions(+) diff --git a/src/utils/bdates.rs b/src/utils/bdates.rs index 7895be6..7c13646 100644 --- a/src/utils/bdates.rs +++ b/src/utils/bdates.rs @@ -148,6 +148,100 @@ enum GroupKey { Yearly(i32), // Group by year } +/// Represents a collection of business 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). +/// +/// Business days are typically Monday to Friday. Weekend dates are skipped or +/// adjusted depending on the frequency rules. +/// +/// # Examples +/// +/// **1. Using `new` (Start and End Date):** +/// +/// ```rust +/// use chrono::NaiveDate; +/// use std::error::Error; +/// # use bdates::{BDatesList, BDateFreq}; // Replace bdates with your actual crate/module name +/// +/// # fn main() -> Result<(), Box> { +/// let start_date = "2023-11-01".to_string(); // Wednesday +/// let end_date = "2023-11-07".to_string(); // Tuesday +/// let freq = BDateFreq::Daily; +/// +/// let bdates = BDatesList::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, 6).unwrap(), // Mon +/// NaiveDate::from_ymd_opt(2023, 11, 7).unwrap(), // Tue +/// ]; +/// +/// assert_eq!(bdates.list()?, expected_dates); +/// assert_eq!(bdates.count()?, 5); +/// # Ok(()) +/// # } +/// ``` +/// +/// **2. Using `from_n_periods` (Start Date and Count):** +/// +/// ```rust +/// use chrono::NaiveDate; +/// use std::error::Error; +/// # use bdates::{BDatesList, BDateFreq}; // Replace bdates with your actual crate/module name +/// +/// # fn main() -> Result<(), Box> { +/// let start_date = "2024-02-28".to_string(); // Wednesday +/// let freq = BDateFreq::WeeklyFriday; +/// let n_periods = 3; +/// +/// let bdates = BDatesList::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!(bdates.list()?, expected_dates); +/// assert_eq!(bdates.count()?, 3); +/// assert_eq!(bdates.start_date_str(), "2024-02-28"); // Keeps original start string +/// assert_eq!(bdates.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 bdates::{BDatesList, BDateFreq}; // Replace bdates with your actual crate/module name +/// +/// # fn main() -> Result<(), Box> { +/// let start_date = "2023-11-20".to_string(); // Mon, Week 47 +/// let end_date = "2023-12-08".to_string(); // Fri, Week 49 +/// let freq = BDateFreq::WeeklyMonday; +/// +/// let bdates = BDatesList::new(start_date, end_date, freq); +/// +/// // Mondays in range: Nov 20, Nov 27, Dec 4 +/// let groups = bdates.groups()?; +/// +/// assert_eq!(groups.len(), 3); // One group per week containing a Monday +/// assert_eq!(groups[0], vec![NaiveDate::from_ymd_opt(2023, 11, 20).unwrap()]); // Week 47 +/// assert_eq!(groups[1], vec![NaiveDate::from_ymd_opt(2023, 11, 27).unwrap()]); // Week 48 +/// assert_eq!(groups[2], vec![NaiveDate::from_ymd_opt(2023, 12, 4).unwrap()]); // Week 49 +/// # Ok(()) +/// # } +/// ``` impl BDatesList { /// Creates a new `BDatesList` instance defined by a start and end date. /// From 1588f05b9de97352507a03c0799de4e10f4dea1d Mon Sep 17 00:00:00 2001 From: Palash Tyagi <23239946+Magnus167@users.noreply.github.com> Date: Sun, 20 Apr 2025 05:24:13 +0100 Subject: [PATCH 3/4] added docstring for bdatesGenerator --- src/utils/bdates.rs | 56 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 56 insertions(+) diff --git a/src/utils/bdates.rs b/src/utils/bdates.rs index 7c13646..ae76476 100644 --- a/src/utils/bdates.rs +++ b/src/utils/bdates.rs @@ -441,6 +441,62 @@ impl BDatesList { /// 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:** +/// +/// ```rust +/// use chrono::NaiveDate; +/// use std::error::Error; +/// # use bdates::{BDatesGenerator, BDateFreq}; // Replace bdates with your actual crate/module name +/// +/// # fn main() -> Result<(), Box> { +/// let start = NaiveDate::from_ymd_opt(2023, 12, 28).unwrap(); // Thursday +/// let freq = BDateFreq::MonthEnd; +/// let n_periods = 4; // Dec '23, Jan '24, Feb '24, Mar '24 +/// +/// let mut generator = BDatesGenerator::new(start, freq, n_periods)?; +/// +/// // First month-end on or after 2023-12-28 is 2023-12-29 +/// assert_eq!(generator.next(), Some(NaiveDate::from_ymd_opt(2023, 12, 29).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, 29).unwrap())); // Mar 31 is Sun +/// assert_eq!(generator.next(), None); // Exhausted +/// # Ok(()) +/// # } +/// ``` +/// +/// **2. Collecting into a Vec:** +/// +/// ```rust +/// use chrono::NaiveDate; +/// use std::error::Error; +/// # use bdates::{BDatesGenerator, BDateFreq}; // Replace bdates with your actual crate/module name +/// +/// # fn main() -> Result<(), Box> { +/// let start = NaiveDate::from_ymd_opt(2024, 4, 29).unwrap(); // Monday +/// let freq = BDateFreq::Daily; +/// let n_periods = 5; +/// +/// let generator = BDatesGenerator::new(start, freq, n_periods)?; +/// let dates: Vec = 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(()) +/// # } +/// ``` #[derive(Debug, Clone)] pub struct BDatesGenerator { freq: BDateFreq, From 0bfadd5f970bf96eeb39beb78770236c5679f38e Mon Sep 17 00:00:00 2001 From: Palash Tyagi <23239946+Magnus167@users.noreply.github.com> Date: Sun, 20 Apr 2025 05:54:09 +0100 Subject: [PATCH 4/4] added modifications to use generator in dates util --- src/utils/dates.rs | 2606 ++++++++++++++++++++++++++++++-------------- 1 file changed, 1769 insertions(+), 837 deletions(-) diff --git a/src/utils/dates.rs b/src/utils/dates.rs index a104a5b..779b469 100644 --- a/src/utils/dates.rs +++ b/src/utils/dates.rs @@ -3,99 +3,61 @@ use std::collections::HashMap; use std::error::Error; use std::hash::Hash; use std::result::Result; +use std::str::FromStr; -// --- Enums --- +// --- Core Enums --- /// Represents the frequency at which calendar dates should be generated. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum DateFreq { - Daily, // Every single day - WeeklyMonday, // The Monday of each week - WeeklyFriday, // The Friday of each week - MonthStart, // The 1st day of each month - MonthEnd, // The actual last day of each month - QuarterStart, // The 1st day of each quarter (Jan 1, Apr 1, Jul 1, Oct 1) - QuarterEnd, // The actual last day of each quarter (Mar 31, Jun 30, Sep 30, Dec 31) - YearStart, // January 1st of each year - YearEnd, // December 31st of each year + 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 date in a group's period. - End, // Indicates picking the last date in a group's period. + Start, // Indicates picking the first calendar date in a group's period. + End, // Indicates picking the last calendar day in a group's period. } -// Helper enum for grouping dates (Internal) -#[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 -} - -// --- DateFreq Implementation --- - impl DateFreq { /// Attempts to parse a frequency string into a `DateFreq` enum. - /// Supports various frequency codes and common aliases. /// - /// | Code | Alias | Description | - /// |------|---------|-------------------------| - /// | D | | Daily | - /// | W | WS | Weekly Monday | - /// | WF | | Weekly Friday | - /// | M | MS | Month Start (1st) | - /// | ME | | Month End (actual last) | - /// | Q | QS | Quarter Start (1st) | - /// | QE | | Quarter End (actual last)| - /// | Y | A, AS, YS | Year Start (Jan 1st) | - /// | YE | AE | Year End (Dec 31st) | + /// This is a convenience wrapper around `from_str`. /// /// # Arguments /// - /// * `freq` - The frequency string slice (e.g., "D", "W", "ME"). + /// * `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_str(freq: &str) -> Result> { - let r = match freq { - "D" => DateFreq::Daily, - "W" | "WS" => DateFreq::WeeklyMonday, - "WF" => DateFreq::WeeklyFriday, - "M" | "MS" => DateFreq::MonthStart, - "ME" => DateFreq::MonthEnd, - "Q" | "QS" => DateFreq::QuarterStart, - "QE" => DateFreq::QuarterEnd, - "Y" | "A" | "AS" | "YS" => DateFreq::YearStart, - "YE" | "AE" => DateFreq::YearEnd, - _ => return Err(format!("Invalid frequency specified: {}", freq).into()), - }; - Ok(r) - } - - /// Attempts to parse a frequency string into a `DateFreq` enum. - /// Convenience wrapper around `from_str`. pub fn from_string(freq: String) -> Result> { - Self::from_str(&freq) + 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::WeeklyFriday => "WF", DateFreq::MonthStart => "M", - DateFreq::MonthEnd => "ME", DateFreq::QuarterStart => "Q", - DateFreq::QuarterEnd => "QE", DateFreq::YearStart => "Y", + DateFreq::MonthEnd => "ME", + DateFreq::QuarterEnd => "QE", + DateFreq::WeeklyFriday => "WF", DateFreq::YearEnd => "YE", }; r.to_string() @@ -118,11 +80,55 @@ impl DateFreq { } } +// Implement FromStr for DateFreq to allow parsing directly using `parse()` +impl FromStr for DateFreq { + type Err = Box; + + /// 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 { + 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) + } +} + // --- DatesList Struct --- /// Represents a list of calendar dates generated between a start and end date -/// at a specified frequency. Includes all days (weekends, etc.). -/// Provides methods to retrieve the full list, count, or dates grouped by period. +/// 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, @@ -130,16 +136,119 @@ pub struct DatesList { freq: DateFreq, } -// --- DatesList Implementation --- +// 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 dates::{DatesList, DateFreq}; // Assuming the crate/module is named 'dates' +/// +/// # fn main() -> Result<(), Box> { +/// 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 dates::{DatesList, DateFreq}; +/// +/// # fn main() -> Result<(), Box> { +/// 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 dates::{DatesList, DateFreq}; +/// +/// # fn main() -> Result<(), Box> { +/// 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. + /// 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 (`DateFreq`) for generating dates. + /// * `freq` - The frequency for generating dates. pub fn new(start_date_str: String, end_date_str: String, freq: DateFreq) -> Self { DatesList { start_date_str, @@ -148,7 +257,53 @@ impl DatesList { } } - /// Returns the flat list of calendar dates within the specified range and frequency. + /// 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> { + 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 = 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. /// @@ -156,39 +311,35 @@ impl DatesList { /// /// Returns an error if the start or end date strings cannot be parsed. pub fn list(&self) -> Result, Box> { + // 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 calendar dates within the specified range and frequency. + /// 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. + /// Returns an error if the start or end date strings cannot be parsed (as it + /// calls `list` internally). pub fn count(&self) -> Result> { self.list().map(|list| list.len()) } - /// Returns a list of date lists, where each inner list contains the dates - /// generated by `self.list()` that belong to the same period (determined by `self.freq`). + /// 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/year. - /// - For `Monthly` frequencies, grouping is by month/year. - /// - For `Quarterly` frequencies, grouping is by quarter/year. - /// - For `Yearly` frequencies, grouping is by year. + /// 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>, Box> { let dates = self.list()?; - if dates.is_empty() { - return Ok(Vec::new()); - } - let mut groups: HashMap> = HashMap::new(); for date in dates { @@ -210,13 +361,18 @@ impl DatesList { } let mut sorted_groups: Vec<(GroupKey, Vec)> = groups.into_iter().collect(); - sorted_groups.sort_by_key(|(k, _)| *k); + 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) } /// 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::parse_from_str(&self.start_date_str, "%Y-%m-%d").map_err(|e| e.into()) } @@ -227,6 +383,10 @@ impl DatesList { } /// 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::parse_from_str(&self.end_date_str, "%Y-%m-%d").map_err(|e| e.into()) } @@ -236,7 +396,7 @@ impl DatesList { &self.end_date_str } - /// Returns the frequency enum (`DateFreq`). + /// Returns the frequency enum. pub fn freq(&self) -> DateFreq { self.freq } @@ -247,9 +407,165 @@ impl DatesList { } } -// --- Internal Helper Functions --- +// --- Dates Generator (Iterator) --- -/// Generates the flat list of calendar dates for the given range and frequency. +/// 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; +/// # // Replace 'your_crate_name' with the actual name if this is in a library +/// # use dates::{DatesGenerator, DateFreq}; +/// +/// # fn main() -> Result<(), Box> { +/// 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 dates::{DatesGenerator, DateFreq}; +/// +/// # fn main() -> Result<(), Box> { +/// 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 = 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 dates::{DatesGenerator, DateFreq}; +/// +/// # fn main() -> Result<(), Box> { +/// 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 { + freq: DateFreq, + periods_remaining: usize, + // Stores the *next* date to be yielded by the iterator. + next_date_candidate: Option, +} + +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> { + 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 { + 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 --- + +/// Generates the flat list of dates for the given range and frequency. +/// Assumes the `collect_*` functions return sorted dates. fn get_dates_list_with_freq( start_date_str: &str, end_date_str: &str, @@ -262,62 +578,73 @@ fn get_dates_list_with_freq( return Ok(Vec::new()); } - let mut dates = match freq { - DateFreq::Daily => collect_calendar_daily(start_date, end_date), - DateFreq::WeeklyMonday => collect_calendar_weekly(start_date, end_date, Weekday::Mon), - DateFreq::WeeklyFriday => collect_calendar_weekly(start_date, end_date, Weekday::Fri), - DateFreq::MonthStart => collect_calendar_monthly(start_date, end_date, true), - DateFreq::MonthEnd => collect_calendar_monthly(start_date, end_date, false), - DateFreq::QuarterStart => collect_calendar_quarterly(start_date, end_date, true), - DateFreq::QuarterEnd => collect_calendar_quarterly(start_date, end_date, false), - DateFreq::YearStart => collect_calendar_yearly(start_date, end_date, true), - DateFreq::YearEnd => collect_calendar_yearly(start_date, end_date, false), + 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)?, }; - // Ensure the final list is sorted (most collectors produce sorted, but good practice). - dates.sort_unstable(); // Slightly faster for non-pathological cases - + // The collect_* functions should now generate sorted dates directly. Ok(dates) } -// --- Internal Date Collection Logic --- +/* ---------------------- Low-Level Date Collection Functions (Internal) ---------------------- */ +// These functions generate dates within a *range* [start_date, end_date] -fn collect_calendar_daily(start_date: NaiveDate, end_date: NaiveDate) -> Vec { +/// Returns all calendar days day-by-day within the range. +fn collect_daily( + start_date: NaiveDate, + end_date: NaiveDate, +) -> Result, Box> { let mut result = Vec::new(); let mut current = start_date; while current <= end_date { result.push(current); - match current.succ_opt() { - Some(next_day) => current = next_day, - None => break, // Avoid panic on date overflow near max date - } + current = current + .succ_opt() + .ok_or("Date overflow near end of supported range")?; } - result + Ok(result) } -fn collect_calendar_weekly( +/// Returns the specified `target_weekday` in each week within the range. +fn collect_weekly( start_date: NaiveDate, end_date: NaiveDate, target_weekday: Weekday, -) -> Vec { +) -> Result, Box> { let mut result = Vec::new(); - let mut current = move_to_weekday_on_or_after(start_date, target_weekday); + let mut current = move_to_day_of_week_on_or_after(start_date, target_weekday)?; while current <= end_date { result.push(current); - match current.checked_add_signed(Duration::days(7)) { - Some(next_week_day) => current = next_week_day, - None => break, // Avoid panic on date overflow - } + current = current + .checked_add_signed(Duration::days(7)) + .ok_or("Date overflow adding 7 days")?; } - result + Ok(result) } -fn collect_calendar_monthly( +/// 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, -) -> Vec { +) -> Result, Box> { let mut result = Vec::new(); let mut year = start_date.year(); let mut month = start_date.month(); @@ -326,280 +653,436 @@ fn collect_calendar_monthly( |(yr, mo): (i32, u32)| -> (i32, u32) { if mo == 12 { (yr + 1, 1) } else { (yr, mo + 1) } }; loop { - let candidate_res = if want_first_day { - NaiveDate::from_ymd_opt(year, month, 1) + let candidate = if want_first_day { + first_day_of_month(year, month)? } else { - days_in_month(year, month).and_then(|day| NaiveDate::from_ymd_opt(year, month, day)) - }; - - let candidate = match candidate_res { - Some(date) => date, - None => { - // Should not happen with valid year/month logic, but break defensively - eprintln!("Warning: Invalid date generated for {}-{}", year, month); - break; - } + 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; - if year > end_date.year() { - break; - } - if year == end_date.year() && month > end_date.month() { - break; + // 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()); } } - result + + Ok(result) } -fn collect_calendar_quarterly( +/// 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, -) -> Vec { +) -> Result, Box> { let mut result = Vec::new(); let mut year = start_date.year(); let mut q = month_to_quarter(start_date.month()); loop { - let candidate = match if want_first_day { - first_day_of_quarter(year, q) + let candidate = if want_first_day { + first_day_of_quarter(year, q)? } else { - last_day_of_quarter(year, q) - } { - Ok(date) => date, - Err(_) => { - // Handle potential panic from helpers if q is invalid - eprintln!( - "Warning: Invalid date generated for quarter {}-Q{}", - year, q - ); - break; - } + 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; } - - // Check if the *start* of the *next* quarter is already past the end date - // to potentially break earlier. - match first_day_of_quarter(year, q) { - Ok(next_q_start) => { - if next_q_start > end_date && want_first_day { - break; - } // If we want start, no need to check further - if next_q_start > end_date && !want_first_day && candidate < start_date { - break; - } // If we want end and haven't found one yet, no need to check further - } - Err(_) => break, // Invalid next quarter, stop - } - - // Basic loop guard + // Safety check if year > end_date.year() + 2 { - eprintln!("Warning: Quarter loop seems excessive, breaking."); - break; + return Err("Loop seems to exceed reasonable year range in collect_quarterly".into()); } } - result + + Ok(result) } -fn collect_calendar_yearly( +/// 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, -) -> Vec { +) -> Result, Box> { let mut result = Vec::new(); let mut year = start_date.year(); while year <= end_date.year() { - let candidate_res = if want_first_day { - NaiveDate::from_ymd_opt(year, 1, 1) + let candidate = if want_first_day { + first_day_of_year(year)? } else { - NaiveDate::from_ymd_opt(year, 12, 31) + last_day_of_year(year)? }; - let candidate = match candidate_res { - Some(date) => date, - None => { - // Should only happen near chrono::MINYEAR/MAXYEAR - eprintln!("Warning: Invalid date generated for year {}", year); - year += 1; // Try next year - continue; - } - }; - - if candidate > end_date { - break; - } // Candidate past range end - if candidate >= start_date { + 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 += 1; + year = year.checked_add(1).ok_or("Year overflow")?; } - result + Ok(result) } -// --- Internal Date Utility Functions --- +/* ---------------------- Core Date Utility Functions (Internal) ---------------------- */ /// 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_weekday_on_or_after(date: NaiveDate, target: Weekday) -> NaiveDate { - let current_weekday_num = date.weekday().num_days_from_monday(); - let target_weekday_num = target.num_days_from_monday(); - let days_forward = (target_weekday_num + 7 - current_weekday_num) % 7; - // Use checked_add for safety near date limits, though expect is common practice here - date.checked_add_signed(Duration::days(days_forward as i64)) - .expect("Date calculation overflow near MAX/MIN date") +fn move_to_day_of_week_on_or_after( + date: NaiveDate, + target: Weekday, +) -> Result> { + let mut current = date; + while current.weekday() != target { + current = current + .succ_opt() + .ok_or("Date overflow moving to next weekday")?; + } + Ok(current) } -/// Returns the number of days in a given month and year. Returns None if month is invalid. -fn days_in_month(year: i32, month: u32) -> Option { +/// Return the first calendar day of the given (year, month). +fn first_day_of_month(year: i32, month: u32) -> Result> { if !(1..=12).contains(&month) { - return None; // Explicitly handle invalid months + 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> { + if !(1..=12).contains(&month) { + return Err(format!("Invalid month: {}", month).into()); } let (ny, nm) = if month == 12 { - (year.checked_add(1)?, 1) + ( + year.checked_add(1) + .ok_or("Year overflow calculating next month")?, + 1, + ) } else { (year, month + 1) }; - let first_of_next = NaiveDate::from_ymd_opt(ny, nm, 1)?; - let last_of_this = first_of_next.pred_opt()?; - Some(last_of_this.day()) + // 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> { + // 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. +/// Panics if month is invalid (should not happen with valid NaiveDate). fn month_to_quarter(m: u32) -> u32 { - assert!((1..=12).contains(&m), "Invalid month: {}", m); - (m - 1) / 3 + 1 + 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 (year, quarter). -/// Returns Err if quarter is invalid or date calculation fails. -fn first_day_of_quarter(year: i32, quarter: u32) -> Result { - let month = match quarter { - 1 => 1, - 2 => 4, - 3 => 7, - 4 => 10, - _ => return Err("Invalid quarter"), - }; - NaiveDate::from_ymd_opt(year, month, 1).ok_or("Invalid date from quarter") +/// Returns the 1st day of the month that starts a given quarter. +fn quarter_start_month(quarter: u32) -> Result> { + 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 + } } -/// Returns the *actual* last calendar day in the given (year, quarter). -/// Returns Err if quarter is invalid or date calculation fails. -fn last_day_of_quarter(year: i32, quarter: u32) -> Result { - let last_month_in_quarter = match quarter { - 1 => 3, - 2 => 6, - 3 => 9, - 4 => 12, - _ => return Err("Invalid quarter"), - }; - let last_day = - days_in_month(year, last_month_in_quarter).ok_or("Invalid month for quarter end")?; - NaiveDate::from_ymd_opt(year, last_month_in_quarter, last_day) - .ok_or("Invalid date for quarter end") +/// Return the first calendar day in the given (year, quarter). +fn first_day_of_quarter(year: i32, quarter: u32) -> Result> { + // Propagate error from quarter_start_month + let month = quarter_start_month(quarter)?; + first_day_of_month(year, month) } -// --- Unit Tests --- +/// Returns the last day of the month that ends a given quarter. +fn quarter_end_month(quarter: u32) -> Result> { + 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> { + // 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::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::from_ymd_opt(year, 12, 31) + .ok_or_else(|| format!("Invalid year for Dec 31st: {}", year).into()) +} + +// --- Generator Helper Functions --- + +/// Finds the *first* valid date according to the frequency, +/// starting the search *on or after* the given `start_date`. +fn find_first_date_on_or_after( + start_date: NaiveDate, + freq: DateFreq, +) -> Result> { + 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 => { + let mut candidate = first_day_of_month(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 = first_day_of_month(next_y, next_m)?; + } + Ok(candidate) + } + DateFreq::MonthEnd => { + let mut candidate = last_day_of_month(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 = last_day_of_month(next_y, next_m)?; + } + Ok(candidate) + } + DateFreq::QuarterStart => { + let current_q = month_to_quarter(start_date.month()); + let mut candidate = first_day_of_quarter(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 = first_day_of_quarter(next_y, next_q)?; + } + Ok(candidate) + } + DateFreq::QuarterEnd => { + let current_q = month_to_quarter(start_date.month()); + let mut candidate = last_day_of_quarter(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 = last_day_of_quarter(next_y, next_q)?; + } + Ok(candidate) + } + DateFreq::YearStart => { + let mut candidate = first_day_of_year(start_date.year())?; + if candidate < start_date { + candidate = + first_day_of_year(start_date.year().checked_add(1).ok_or("Year overflow")?)?; + } + Ok(candidate) + } + DateFreq::YearEnd => { + let mut candidate = last_day_of_year(start_date.year())?; + if candidate < start_date { + candidate = + last_day_of_year(start_date.year().checked_add(1).ok_or("Year overflow")?)?; + } + 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). +fn find_next_date(current_date: NaiveDate, freq: DateFreq) -> Result> { + 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 => { + 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) + }; + first_day_of_month(next_y, next_m) + } + DateFreq::MonthEnd => { + 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) + }; + last_day_of_month(next_y, next_m) + } + DateFreq::QuarterStart => { + let current_q = month_to_quarter(current_date.month()); + 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) + }; + first_day_of_quarter(next_y, next_q) + } + DateFreq::QuarterEnd => { + let current_q = month_to_quarter(current_date.month()); + 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) + }; + last_day_of_quarter(next_y, next_q) + } + DateFreq::YearStart => { + first_day_of_year(current_date.year().checked_add(1).ok_or("Year overflow")?) + } + DateFreq::YearEnd => { + last_day_of_year(current_date.year().checked_add(1).ok_or("Year overflow")?) + } + } +} + +// --- Tests --- #[cfg(test)] mod tests { - use super::*; // Import everything from the parent module - use chrono::NaiveDate; + use super::*; + use chrono::{Duration, NaiveDate, Weekday}; // Make sure Duration is imported - // Helper function to create NaiveDate instances easily in tests - fn d(year: i32, month: u32, day: u32) -> NaiveDate { - NaiveDate::from_ymd_opt(year, month, day).unwrap() + // 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") } - // --- Tests for DateFreq --- + // --- DateFreq Tests --- #[test] - fn test_date_freq_from_str_valid() { - assert_eq!(DateFreq::from_str("D").unwrap(), DateFreq::Daily); - assert_eq!(DateFreq::from_str("W").unwrap(), DateFreq::WeeklyMonday); - assert_eq!(DateFreq::from_str("WS").unwrap(), DateFreq::WeeklyMonday); - assert_eq!(DateFreq::from_str("WF").unwrap(), DateFreq::WeeklyFriday); - assert_eq!(DateFreq::from_str("M").unwrap(), DateFreq::MonthStart); - assert_eq!(DateFreq::from_str("MS").unwrap(), DateFreq::MonthStart); - assert_eq!(DateFreq::from_str("ME").unwrap(), DateFreq::MonthEnd); - assert_eq!(DateFreq::from_str("Q").unwrap(), DateFreq::QuarterStart); - assert_eq!(DateFreq::from_str("QS").unwrap(), DateFreq::QuarterStart); - assert_eq!(DateFreq::from_str("QE").unwrap(), DateFreq::QuarterEnd); - assert_eq!(DateFreq::from_str("Y").unwrap(), DateFreq::YearStart); - assert_eq!(DateFreq::from_str("A").unwrap(), DateFreq::YearStart); - assert_eq!(DateFreq::from_str("AS").unwrap(), DateFreq::YearStart); - assert_eq!(DateFreq::from_str("YS").unwrap(), DateFreq::YearStart); - assert_eq!(DateFreq::from_str("YE").unwrap(), DateFreq::YearEnd); - assert_eq!(DateFreq::from_str("AE").unwrap(), DateFreq::YearEnd); + fn test_datefreq_from_str() -> Result<(), Box> { + assert_eq!(DateFreq::from_str("D")?, DateFreq::Daily); + assert_eq!("D".parse::()?, 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::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::().is_err()); + let err = DateFreq::from_str("INVALID").unwrap_err(); + assert_eq!(err.to_string(), "Invalid frequency specified: INVALID"); + + Ok(()) } #[test] - fn test_date_freq_from_string_valid() { - // Test the wrapper function - assert_eq!( - DateFreq::from_string("D".to_string()).unwrap(), - DateFreq::Daily - ); - assert_eq!( - DateFreq::from_string("YE".to_string()).unwrap(), - DateFreq::YearEnd - ); - } - - #[test] - fn test_date_freq_from_str_invalid() { - assert!(DateFreq::from_str("X").is_err()); - assert!(DateFreq::from_str(" Monthly").is_err()); - assert!(DateFreq::from_str("").is_err()); - } - - #[test] - fn test_date_freq_from_string_invalid() { - assert!(DateFreq::from_string("Invalid Freq".to_string()).is_err()); - } - - #[test] - fn test_date_freq_to_string() { + fn test_datefreq_to_string() { assert_eq!(DateFreq::Daily.to_string(), "D"); assert_eq!(DateFreq::WeeklyMonday.to_string(), "W"); - assert_eq!(DateFreq::WeeklyFriday.to_string(), "WF"); assert_eq!(DateFreq::MonthStart.to_string(), "M"); - assert_eq!(DateFreq::MonthEnd.to_string(), "ME"); assert_eq!(DateFreq::QuarterStart.to_string(), "Q"); - assert_eq!(DateFreq::QuarterEnd.to_string(), "QE"); 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_date_freq_agg_type() { + fn test_datefreq_from_string() -> Result<(), Box> { + 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); @@ -612,518 +1095,473 @@ mod tests { assert_eq!(DateFreq::YearEnd.agg_type(), AggregationType::End); } - // --- Tests for DatesList Accessors --- + // --- DatesList Property Tests --- #[test] - fn test_dates_list_accessors() { - let start = "2024-01-15"; - let end = "2024-02-20"; - let freq = DateFreq::WeeklyMonday; - let dl = DatesList::new(start.to_string(), end.to_string(), freq); + fn test_dates_list_properties_new() -> Result<(), Box> { + 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!(dl.start_date_str(), start); - assert_eq!(dl.end_date_str(), end); - assert_eq!(dl.start_date().unwrap(), d(2024, 1, 15)); - assert_eq!(dl.end_date().unwrap(), d(2024, 2, 20)); - assert_eq!(dl.freq(), freq); - assert_eq!(dl.freq_str(), "W"); // Canonical string for WeeklyMonday + 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_invalid_dates() { - let dl_bad_start = DatesList::new( - "2024-13-01".to_string(), - "2024-01-10".to_string(), + fn test_dates_list_properties_from_n_periods() -> Result<(), Box> { + 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!(dl_bad_start.start_date().is_err()); - assert!(dl_bad_start.list().is_err()); // list() should propagate parse error - assert!(dl_bad_start.count().is_err()); - assert!(dl_bad_start.groups().is_err()); + 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 dl_bad_end = DatesList::new( - "2024-01-01".to_string(), + let list_end_invalid = DatesList::new( + "2023-01-01".to_string(), "invalid-date".to_string(), DateFreq::Daily, ); - assert!(dl_bad_end.end_date().is_err()); - assert!(dl_bad_end.list().is_err()); - assert!(dl_bad_end.count().is_err()); - assert!(dl_bad_end.groups().is_err()); + 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()); } - // --- Tests for DatesList::list() and DatesList::count() --- + // --- DatesList Core Logic Tests (list, count) --- #[test] - fn test_dates_list_empty_range() { - let dl = DatesList::new( - "2024-01-10".to_string(), - "2024-01-09".to_string(), + fn test_dates_list_daily_list() -> Result<(), Box> { + let dates_list = DatesList::new( + "2023-11-01".to_string(), // Wed + "2023-11-05".to_string(), // Sun DateFreq::Daily, ); - assert_eq!(dl.list().unwrap(), Vec::new()); - assert_eq!(dl.count().unwrap(), 0); - assert_eq!(dl.groups().unwrap(), Vec::>::new()); // Ensure groups also handles empty + 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_single_day_range() { - let dl = DatesList::new( - "2024-01-10".to_string(), - "2024-01-10".to_string(), + fn test_dates_list_weekly_monday_list() -> Result<(), Box> { + 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> { + 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> { + 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> { + 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> { + 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> { + 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> { + 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> { + 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> { + let dates_list = DatesList::new( + "2023-12-31".to_string(), + "2023-01-01".to_string(), // End date before start date DateFreq::Daily, ); - assert_eq!(dl.list().unwrap(), vec![d(2024, 1, 10)]); - assert_eq!(dl.count().unwrap(), 1); + 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> { + 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_daily() { - let dl = DatesList::new( - "2024-03-29".to_string(), - "2024-04-02".to_string(), + fn test_dates_list_groups_daily() -> Result<(), Box> { + let dates_list = DatesList::new( + "2023-11-01".to_string(), // Wed + "2023-11-03".to_string(), // Fri DateFreq::Daily, ); - let expected = vec![ - d(2024, 3, 29), // Fri - d(2024, 3, 30), // Sat - d(2024, 3, 31), // Sun - d(2024, 4, 1), // Mon - d(2024, 4, 2), // Tue - ]; - assert_eq!(dl.list().unwrap(), expected); - assert_eq!(dl.count().unwrap(), 5); + 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_weekly_monday() { - // Start before first Mon, end after last Mon - let dl1 = DatesList::new( - "2024-01-02".to_string(), - "2024-01-20".to_string(), - DateFreq::WeeklyMonday, - ); - assert_eq!(dl1.list().unwrap(), vec![d(2024, 1, 8), d(2024, 1, 15)]); - assert_eq!(dl1.count().unwrap(), 2); - - // Start on a Mon, end on a Mon - let dl2 = DatesList::new( - "2024-01-08".to_string(), - "2024-01-22".to_string(), - DateFreq::WeeklyMonday, - ); - assert_eq!( - dl2.list().unwrap(), - vec![d(2024, 1, 8), d(2024, 1, 15), d(2024, 1, 22)] - ); - assert_eq!(dl2.count().unwrap(), 3); - - // Across year boundary - let dl3 = DatesList::new( - "2023-12-28".to_string(), - "2024-01-10".to_string(), - DateFreq::WeeklyMonday, - ); - assert_eq!(dl3.list().unwrap(), vec![d(2024, 1, 1), d(2024, 1, 8)]); - assert_eq!(dl3.count().unwrap(), 2); - } - - #[test] - fn test_dates_list_weekly_friday() { - // Start before first Fri, end after last Fri - let dl1 = DatesList::new( - "2024-01-01".to_string(), - "2024-01-20".to_string(), + fn test_dates_list_groups_weekly_friday() -> Result<(), Box> { + 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, ); - assert_eq!( - dl1.list().unwrap(), - vec![d(2024, 1, 5), d(2024, 1, 12), d(2024, 1, 19)] - ); - assert_eq!(dl1.count().unwrap(), 3); - - // Start on a Fri, end on a Fri - let dl2 = DatesList::new( - "2024-01-12".to_string(), - "2024-01-26".to_string(), - DateFreq::WeeklyFriday, - ); - assert_eq!( - dl2.list().unwrap(), - vec![d(2024, 1, 12), d(2024, 1, 19), d(2024, 1, 26)] - ); - assert_eq!(dl2.count().unwrap(), 3); - - // Across year boundary - let dl3 = DatesList::new( - "2023-12-25".to_string(), - "2024-01-15".to_string(), - DateFreq::WeeklyFriday, - ); - assert_eq!( - dl3.list().unwrap(), - vec![d(2023, 12, 29), d(2024, 1, 5), d(2024, 1, 12)] - ); - assert_eq!(dl3.count().unwrap(), 3); + 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_month_start() { - // Basic range - let dl1 = DatesList::new( - "2024-01-15".to_string(), - "2024-04-10".to_string(), - DateFreq::MonthStart, - ); - assert_eq!( - dl1.list().unwrap(), - vec![d(2024, 2, 1), d(2024, 3, 1), d(2024, 4, 1)] - ); - assert_eq!(dl1.count().unwrap(), 3); - - // Start exactly on MonthStart - let dl2 = DatesList::new( - "2024-02-01".to_string(), - "2024-03-15".to_string(), - DateFreq::MonthStart, - ); - assert_eq!(dl2.list().unwrap(), vec![d(2024, 2, 1), d(2024, 3, 1)]); - assert_eq!(dl2.count().unwrap(), 2); - - // Across year boundary - let dl3 = DatesList::new( - "2023-11-20".to_string(), - "2024-02-10".to_string(), - DateFreq::MonthStart, - ); - assert_eq!( - dl3.list().unwrap(), - vec![d(2023, 12, 1), d(2024, 1, 1), d(2024, 2, 1)] - ); - assert_eq!(dl3.count().unwrap(), 3); - } - - #[test] - fn test_dates_list_month_end() { - // Basic range including leap year - let dl1 = DatesList::new( - "2024-01-15".to_string(), - "2024-04-10".to_string(), - DateFreq::MonthEnd, - ); - assert_eq!( - dl1.list().unwrap(), - vec![d(2024, 1, 31), d(2024, 2, 29), d(2024, 3, 31)] - ); // Feb 29 leap - assert_eq!(dl1.count().unwrap(), 3); - - // Start exactly on MonthEnd - let dl2 = DatesList::new( - "2023-11-30".to_string(), - "2024-01-15".to_string(), - DateFreq::MonthEnd, - ); - assert_eq!(dl2.list().unwrap(), vec![d(2023, 11, 30), d(2023, 12, 31)]); - assert_eq!(dl2.count().unwrap(), 2); - - // Non-leap year Feb - let dl3 = DatesList::new( - "2023-01-20".to_string(), - "2023-03-10".to_string(), - DateFreq::MonthEnd, - ); - assert_eq!(dl3.list().unwrap(), vec![d(2023, 1, 31), d(2023, 2, 28)]); // Feb 28 non-leap - assert_eq!(dl3.count().unwrap(), 2); - } - - #[test] - fn test_dates_list_quarter_start() { - // Basic range - let dl1 = DatesList::new( - "2024-02-15".to_string(), - "2024-08-10".to_string(), + fn test_dates_list_groups_quarterly_start() -> Result<(), Box> { + let dates_list = DatesList::new( + "2023-08-01".to_string(), // Start Q3 + "2024-05-01".to_string(), // Start Q2 DateFreq::QuarterStart, ); - assert_eq!(dl1.list().unwrap(), vec![d(2024, 4, 1), d(2024, 7, 1)]); - assert_eq!(dl1.count().unwrap(), 2); - - // Start exactly on QuarterStart - let dl2 = DatesList::new( - "2024-01-01".to_string(), - "2024-07-01".to_string(), - DateFreq::QuarterStart, - ); - assert_eq!( - dl2.list().unwrap(), - vec![d(2024, 1, 1), d(2024, 4, 1), d(2024, 7, 1)] - ); - assert_eq!(dl2.count().unwrap(), 3); - - // Across year boundary - let dl3 = DatesList::new( - "2023-10-15".to_string(), - "2024-05-10".to_string(), - DateFreq::QuarterStart, - ); - assert_eq!(dl3.list().unwrap(), vec![d(2024, 1, 1), d(2024, 4, 1)]); - assert_eq!(dl3.count().unwrap(), 2); + 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_quarter_end() { - // Basic range - let dl1 = DatesList::new( - "2024-02-15".to_string(), - "2024-10-10".to_string(), - DateFreq::QuarterEnd, - ); - assert_eq!( - dl1.list().unwrap(), - vec![d(2024, 3, 31), d(2024, 6, 30), d(2024, 9, 30)] - ); - assert_eq!(dl1.count().unwrap(), 3); - - // End exactly on QuarterEnd - let dl2 = DatesList::new( - "2024-05-01".to_string(), - "2024-09-30".to_string(), - DateFreq::QuarterEnd, - ); - assert_eq!(dl2.list().unwrap(), vec![d(2024, 6, 30), d(2024, 9, 30)]); - assert_eq!(dl2.count().unwrap(), 2); - - // Across year boundary (includes leap year effect on Mar 31 if applicable) - let dl3 = DatesList::new( - "2023-11-20".to_string(), - "2024-04-10".to_string(), - DateFreq::QuarterEnd, - ); - assert_eq!(dl3.list().unwrap(), vec![d(2023, 12, 31), d(2024, 3, 31)]); - assert_eq!(dl3.count().unwrap(), 2); - } - - #[test] - fn test_dates_list_year_start() { - // Basic range - let dl1 = DatesList::new( - "2023-05-15".to_string(), - "2026-08-10".to_string(), - DateFreq::YearStart, - ); - assert_eq!( - dl1.list().unwrap(), - vec![d(2024, 1, 1), d(2025, 1, 1), d(2026, 1, 1)] - ); - assert_eq!(dl1.count().unwrap(), 3); - - // Start exactly on YearStart - let dl2 = DatesList::new( - "2024-01-01".to_string(), - "2025-01-01".to_string(), - DateFreq::YearStart, - ); - assert_eq!(dl2.list().unwrap(), vec![d(2024, 1, 1), d(2025, 1, 1)]); - assert_eq!(dl2.count().unwrap(), 2); - - // Short range within a year - let dl3 = DatesList::new( - "2024-02-01".to_string(), - "2024-11-30".to_string(), - DateFreq::YearStart, - ); - assert_eq!(dl3.list().unwrap(), Vec::::new()); // No Jan 1st within this range - assert_eq!(dl3.count().unwrap(), 0); - } - - #[test] - fn test_dates_list_year_end() { - // Basic range - let dl1 = DatesList::new( - "2023-05-15".to_string(), - "2026-08-10".to_string(), + fn test_dates_list_groups_yearly_end() -> Result<(), Box> { + let dates_list = DatesList::new( + "2022-01-01".to_string(), // Start 2022 + "2024-03-31".to_string(), // Q1 2024 DateFreq::YearEnd, ); - assert_eq!( - dl1.list().unwrap(), - vec![d(2023, 12, 31), d(2024, 12, 31), d(2025, 12, 31)] - ); - assert_eq!(dl1.count().unwrap(), 3); - - // End exactly on YearEnd - let dl2 = DatesList::new( - "2024-01-01".to_string(), - "2025-12-31".to_string(), - DateFreq::YearEnd, - ); - assert_eq!(dl2.list().unwrap(), vec![d(2024, 12, 31), d(2025, 12, 31)]); - assert_eq!(dl2.count().unwrap(), 2); - - // Short range within a year - let dl3 = DatesList::new( - "2024-02-01".to_string(), - "2024-11-30".to_string(), - DateFreq::YearEnd, - ); - assert_eq!(dl3.list().unwrap(), Vec::::new()); // No Dec 31st within this range - assert_eq!(dl3.count().unwrap(), 0); + 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(()) } - // --- Tests for DatesList::groups() --- - #[test] - fn test_dates_list_groups_daily() { - // Grouping daily just puts each date in its own group - let dl = DatesList::new( - "2024-01-01".to_string(), - "2024-01-03".to_string(), + fn test_dates_list_groups_empty_range() -> Result<(), Box> { + let dates_list = DatesList::new( + "2023-12-31".to_string(), + "2023-01-01".to_string(), // End < Start DateFreq::Daily, ); - let expected = vec![ - vec![d(2024, 1, 1)], - vec![d(2024, 1, 2)], - vec![d(2024, 1, 3)], - ]; - assert_eq!(dl.groups().unwrap(), expected); + let groups = dates_list.groups()?; + assert!(groups.is_empty()); + Ok(()) + } + + // --- Tests for internal helper functions --- + + #[test] + fn test_move_to_day_of_week_on_or_after() -> Result<(), Box> { + 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_dates_list_groups_weekly() { - // Test grouping by ISO week, crossing month/year - // Dates generated will be Mondays within the range - let dl = DatesList::new( - "2023-12-28".to_string(), - "2024-01-16".to_string(), - DateFreq::WeeklyMonday, - ); - // Dates: 2024-01-01 (Week 1 2024), 2024-01-08 (Week 2 2024), 2024-01-15 (Week 3 2024) - let expected = vec![ - vec![d(2024, 1, 1)], // Group for Week 1 2024 - vec![d(2024, 1, 8)], // Group for Week 2 2024 - vec![d(2024, 1, 15)], // Group for Week 3 2024 - ]; - assert_eq!(dl.groups().unwrap(), expected); - - // Test Weekly Friday grouping - let dl_fri = DatesList::new( - "2024-01-01".to_string(), - "2024-01-12".to_string(), - DateFreq::WeeklyFriday, - ); - // Dates: 2024-01-05 (Week 1), 2024-01-12 (Week 2) - let expected_fri = vec![ - vec![d(2024, 1, 5)], // Group for Week 1 2024 - vec![d(2024, 1, 12)], // Group for Week 2 2024 - ]; - assert_eq!(dl_fri.groups().unwrap(), expected_fri); + fn test_first_day_of_month() -> Result<(), Box> { + 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_dates_list_groups_monthly() { - let dl = DatesList::new( - "2023-11-15".to_string(), - "2024-02-15".to_string(), - DateFreq::MonthStart, - ); - // Dates: 2023-12-01, 2024-01-01, 2024-02-01 - // Groups: (2023, 12), (2024, 1), (2024, 2) - let expected = vec![ - vec![d(2023, 12, 1)], // Group 2023-12 - vec![d(2024, 1, 1)], // Group 2024-01 - vec![d(2024, 2, 1)], // Group 2024-02 - ]; - assert_eq!(dl.groups().unwrap(), expected); - - // Test Month End grouping - let dl_me = DatesList::new( - "2024-01-20".to_string(), - "2024-03-10".to_string(), - DateFreq::MonthEnd, - ); - // Dates: 2024-01-31, 2024-02-29 - let expected_me = vec![ - vec![d(2024, 1, 31)], // Group 2024-01 - vec![d(2024, 2, 29)], // Group 2024-02 - ]; - assert_eq!(dl_me.groups().unwrap(), expected_me); + fn test_days_in_month() -> Result<(), Box> { + assert_eq!(days_in_month(2023, 1)?, 31); + assert_eq!(days_in_month(2023, 2)?, 28); + assert_eq!(days_in_month(2024, 2)?, 29); // Leap + assert_eq!(days_in_month(2023, 4)?, 30); + assert_eq!(days_in_month(2023, 12)?, 31); + assert!(days_in_month(2023, 0).is_err()); // Invalid month 0 + assert!(days_in_month(2023, 13).is_err()); // Invalid month 13 + // Test near max date year overflow - Use MAX.year() + assert!(days_in_month(NaiveDate::MAX.year(), 12).is_err()); + Ok(()) } #[test] - fn test_dates_list_groups_quarterly() { - let dl = DatesList::new( - "2023-08-01".to_string(), - "2024-05-01".to_string(), - DateFreq::QuarterStart, - ); - // Dates: 2023-10-01 (Q4), 2024-01-01 (Q1), 2024-04-01 (Q2) - // Groups: (2023, 4), (2024, 1), (2024, 2) - let expected = vec![ - vec![d(2023, 10, 1)], // Group 2023-Q4 - vec![d(2024, 1, 1)], // Group 2024-Q1 - vec![d(2024, 4, 1)], // Group 2024-Q2 - ]; - assert_eq!(dl.groups().unwrap(), expected); - - // Test Quarter End grouping - let dl_qe = DatesList::new( - "2023-11-01".to_string(), - "2024-04-15".to_string(), - DateFreq::QuarterEnd, - ); - // Dates: 2023-12-31 (Q4), 2024-03-31 (Q1) - let expected_qe = vec![ - vec![d(2023, 12, 31)], // Group 2023-Q4 - vec![d(2024, 3, 31)], // Group 2024-Q1 - ]; - assert_eq!(dl_qe.groups().unwrap(), expected_qe); + fn test_last_day_of_month() -> Result<(), Box> { + 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)); + assert!(last_day_of_month(2023, 0).is_err()); // Invalid month 0 + assert!(last_day_of_month(2023, 13).is_err()); // Invalid month 13 + // Test near max date year overflow - use MAX.year() + assert!(last_day_of_month(NaiveDate::MAX.year(), 12).is_err()); + Ok(()) } #[test] - fn test_dates_list_groups_yearly() { - let dl = DatesList::new( - "2022-05-01".to_string(), - "2024-12-31".to_string(), - DateFreq::YearEnd, - ); - // Dates: 2022-12-31, 2023-12-31, 2024-12-31 - // Groups: (2022), (2023), (2024) - let expected = vec![ - vec![d(2022, 12, 31)], // Group 2022 - vec![d(2023, 12, 31)], // Group 2023 - vec![d(2024, 12, 31)], // Group 2024 - ]; - assert_eq!(dl.groups().unwrap(), expected); - - // Test Year Start grouping - let dl_ys = DatesList::new( - "2023-02-01".to_string(), - "2025-01-01".to_string(), - DateFreq::YearStart, - ); - // Dates: 2024-01-01, 2025-01-01 - let expected_ys = vec![ - vec![d(2024, 1, 1)], // Group 2024 - vec![d(2025, 1, 1)], // Group 2025 - ]; - assert_eq!(dl_ys.groups().unwrap(), expected_ys); + 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); } - // --- Tests for Utility Functions (Direct Testing) --- - - #[test] - fn test_days_in_month() { - assert_eq!(days_in_month(2024, 1).unwrap(), 31); // Jan - assert_eq!(days_in_month(2024, 2).unwrap(), 29); // Feb Leap - assert_eq!(days_in_month(2023, 2).unwrap(), 28); // Feb Non-Leap - assert_eq!(days_in_month(2024, 4).unwrap(), 30); // Apr - assert_eq!(days_in_month(2024, 12).unwrap(), 31); // Dec - - // Test invalid months - assert!(days_in_month(2024, 0).is_none()); - assert!(days_in_month(2024, 13).is_none()); - } - - // Panic tests for functions with assertions/panics - #[test] #[should_panic(expected = "Invalid month: 0")] fn test_month_to_quarter_invalid_low() { @@ -1136,123 +1574,617 @@ mod tests { } #[test] - fn test_first_day_of_quarter() { - assert_eq!(first_day_of_quarter(2024, 1).unwrap(), d(2024, 1, 1)); - assert_eq!(first_day_of_quarter(2024, 2).unwrap(), d(2024, 4, 1)); - assert_eq!(first_day_of_quarter(2024, 3).unwrap(), d(2024, 7, 1)); - assert_eq!(first_day_of_quarter(2024, 4).unwrap(), d(2024, 10, 1)); - - assert!(first_day_of_quarter(2024, 0).is_err()); - assert!(first_day_of_quarter(2024, 5).is_err()); + 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_last_day_of_quarter() { - assert_eq!(last_day_of_quarter(2024, 1).unwrap(), d(2024, 3, 31)); - assert_eq!(last_day_of_quarter(2024, 2).unwrap(), d(2024, 6, 30)); - assert_eq!(last_day_of_quarter(2024, 3).unwrap(), d(2024, 9, 30)); - assert_eq!(last_day_of_quarter(2024, 4).unwrap(), d(2024, 12, 31)); - // Check leap year effect indirectly via days_in_month tested elsewhere - assert_eq!(last_day_of_quarter(2023, 1).unwrap(), d(2023, 3, 31)); // Non-leap doesn't affect Q1 end - - assert!(last_day_of_quarter(2024, 0).is_err()); - assert!(last_day_of_quarter(2024, 5).is_err()); + fn test_first_day_of_quarter() -> Result<(), Box> { + 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)); + assert!(first_day_of_quarter(2023, 5).is_err()); // Invalid quarter + Ok(()) } #[test] - fn test_move_to_weekday_on_or_after() { - // Start Mon -> target Mon - assert_eq!( - move_to_weekday_on_or_after(d(2024, 1, 1), Weekday::Mon), - d(2024, 1, 1) - ); - // Start Mon -> target Fri - assert_eq!( - move_to_weekday_on_or_after(d(2024, 1, 1), Weekday::Fri), - d(2024, 1, 5) - ); - // Start Tue -> target Mon - assert_eq!( - move_to_weekday_on_or_after(d(2024, 1, 2), Weekday::Mon), - d(2024, 1, 8) - ); - // Start Sat -> target Mon - assert_eq!( - move_to_weekday_on_or_after(d(2024, 1, 6), Weekday::Mon), - d(2024, 1, 8) - ); - // Start Sun -> target Sun - assert_eq!( - move_to_weekday_on_or_after(d(2024, 1, 7), Weekday::Sun), - d(2024, 1, 7) - ); - // Start Sun -> target Sat (next week) - assert_eq!( - move_to_weekday_on_or_after(d(2024, 1, 7), Weekday::Sat), - d(2024, 1, 13) - ); - } - - // Test potential overflow cases - #[test] - fn test_collect_calendar_near_max_date() { - // Note: NaiveDate::MAX can cause issues with succ_opt/pred_opt in helpers like days_in_month - // Let's test slightly away from the absolute max/min - let end_date = NaiveDate::MAX.pred_opt().unwrap(); // Max - 1 day - let start_date = end_date.pred_opt().unwrap().pred_opt().unwrap(); // Max - 3 days - - let dl = DatesList::new( - start_date.to_string(), - end_date.to_string(), - DateFreq::Daily, - ); - let expected = vec![start_date, start_date.succ_opt().unwrap(), end_date]; - assert_eq!(dl.list().unwrap(), expected); - - // Test weekly near max date - just ensure it doesn't panic - let dl_weekly = DatesList::new( - start_date.to_string(), - end_date.to_string(), - DateFreq::WeeklyMonday, - ); - assert!(dl_weekly.list().is_ok()); - - // Test monthly near max date - let dl_monthly = DatesList::new( - start_date.to_string(), - end_date.to_string(), - DateFreq::MonthEnd, - ); - assert!(dl_monthly.list().is_ok()); + 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_collect_calendar_near_min_date() { - let start_date = NaiveDate::MIN.succ_opt().unwrap(); // Min + 1 day - let end_date = start_date.succ_opt().unwrap().succ_opt().unwrap(); // Min + 3 days + fn test_last_day_of_quarter() -> Result<(), Box> { + 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)); + assert_eq!(last_day_of_quarter(2024, 1)?, date(2024, 3, 31)); // Leap year doesn't affect March end + assert!(last_day_of_quarter(2023, 5).is_err()); // Invalid quarter + // Test overflow propagation - use MAX.year() + assert!(last_day_of_quarter(NaiveDate::MAX.year(), 4).is_err()); + Ok(()) + } - let dl = DatesList::new( - start_date.to_string(), - end_date.to_string(), - DateFreq::Daily, - ); - let expected = vec![start_date, start_date.succ_opt().unwrap(), end_date]; - assert_eq!(dl.list().unwrap(), expected); + #[test] + fn test_first_day_of_year() -> Result<(), Box> { + 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 weekly near min date - let dl_weekly = DatesList::new( - start_date.to_string(), - end_date.to_string(), - DateFreq::WeeklyMonday, - ); - assert!(dl_weekly.list().is_ok()); + #[test] + fn test_last_day_of_year() -> Result<(), Box> { + assert_eq!(last_day_of_year(2023)?, date(2023, 12, 31)); + assert_eq!(last_day_of_year(2024)?, date(2024, 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(NaiveDate::MAX.year())?, NaiveDate::MAX); + Ok(()) + } - // Test monthly near min date - let dl_monthly = DatesList::new( - start_date.to_string(), - end_date.to_string(), - DateFreq::MonthStart, + // Overflow tests for collect_* removed as they were misleading + + // --- Tests for Generator Helper Functions --- + + #[test] + fn test_find_first_date_on_or_after() -> Result<(), Box> { + // Daily + assert_eq!( + find_first_date_on_or_after(date(2023, 11, 8), DateFreq::Daily)?, + date(2023, 11, 8) ); - assert!(dl_monthly.list().is_ok()); + assert_eq!( + find_first_date_on_or_after(date(2023, 11, 11), DateFreq::Daily)?, + date(2023, 11, 11) + ); // Sat -> Sat + + // 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) + ); + assert_eq!( + find_first_date_on_or_after(date(2023, 11, 12), DateFreq::WeeklyMonday)?, + date(2023, 11, 13) + ); // Sun -> Mon + + // 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) + ); + assert_eq!( + find_first_date_on_or_after(date(2023, 10, 1), DateFreq::MonthStart)?, + date(2023, 10, 1) + ); // Oct 1 -> Oct 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) + ); + assert_eq!( + find_first_date_on_or_after(date(2023, 12, 31), DateFreq::MonthEnd)?, + date(2023, 12, 31) + ); // Dec 31 -> Dec 31 + assert_eq!( + find_first_date_on_or_after(date(2024, 2, 15), DateFreq::MonthEnd)?, + date(2024, 2, 29) + ); // Mid Feb (Leap) -> Feb 29 + assert_eq!( + find_first_date_on_or_after(date(2024, 2, 29), DateFreq::MonthEnd)?, + date(2024, 2, 29) + ); // Feb 29 -> Feb 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> { + // 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> { + 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> { + 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> { + 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![ + 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> { + 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> { + 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> { + 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> { + 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> { + 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> { + 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> { + 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> { + 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> { + 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 + + // Check internal state after the call that returned None + // When Some(YE MAX) was returned, periods_remaining became 1. + // The next call enters the match, calls find_next_date (fails -> .ok() is None), + // sets next_date_candidate=None, decrements periods_remaining to 0, returns Some(YE MAX). + // --> NO, the code was: set candidate=find().ok(), THEN decrement. + // Let's revisit Iterator::next logic: + // 1. periods_remaining = 1, next_date_candidate = Some(YE MAX) + // 2. Enter match arm + // 3. find_next_date(YE MAX, YE) -> Err + // 4. self.next_date_candidate = Err.ok() -> None + // 5. self.periods_remaining -= 1 -> becomes 0 + // 6. return Some(YE MAX) <-- This was the bug in my reasoning. It returns the *current* date first. + // State after returning Some(YE MAX): periods_remaining = 0, next_date_candidate = None + // Next call to generator.next(): + // 1. periods_remaining = 0 + // 2. Enter the `_` arm of the match + // 3. self.periods_remaining = 0 (no change) + // 4. self.next_date_candidate = None (no change) + // 5. return None + + // State after the *first* None is returned: + assert_eq!(generator.periods_remaining, 0); // Corrected assertion + 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