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