Add BDateFreq enum and frequency-based date retrieval functions

src/utils/bdates.rs

@@ -0,0 +1,365 @@
+use chrono::{Datelike, Duration, NaiveDate, Weekday};
+use std::error::Error;
+
+#[derive(Debug, Clone, Copy)]
+pub enum BDateFreq {
+    Daily,
+    WeeklyMonday,
+    MonthStart,
+    QuarterStart,
+    YearStart,
+    MonthEnd,
+    QuarterEnd,
+    WeeklyFriday,
+    YearEnd,
+}
+
+#[derive(Debug, Clone, Copy)]
+pub enum AggregationType {
+    Start, // Indicates picking the first date in a group.
+    End,   // Indicates picking the last date in a group.
+}
+
+impl BDateFreq {
+    pub fn from_string(freq: String) -> Result<Self, Box<dyn Error>> {
+        // use from_str to convert the string to a BDateFreq enum
+        Self::from_str(&freq)
+    }
+    pub fn from_str(freq: &str) -> Result<Self, Box<dyn Error>> {
+        match freq {
+            "D" => Ok(BDateFreq::Daily),
+            "W" => Ok(BDateFreq::WeeklyMonday),
+            "M" => Ok(BDateFreq::MonthStart),
+            "Q" => Ok(BDateFreq::QuarterStart),
+            "A" => Ok(BDateFreq::YearStart),
+            "ME" => Ok(BDateFreq::MonthEnd),
+            "QE" => Ok(BDateFreq::QuarterEnd),
+            "WF" => Ok(BDateFreq::WeeklyFriday),
+            "YE" => Ok(BDateFreq::YearEnd),
+            _ => Err("Invalid frequency specified".into()),
+        }
+    }
+    pub fn agg_type(&self) -> AggregationType {
+        match self {
+            BDateFreq::Daily
+            | BDateFreq::WeeklyMonday
+            | BDateFreq::MonthStart
+            | BDateFreq::QuarterStart
+            | BDateFreq::YearStart => AggregationType::Start,
+
+            BDateFreq::WeeklyFriday
+            | BDateFreq::MonthEnd
+            | BDateFreq::QuarterEnd
+            | BDateFreq::YearEnd => AggregationType::End,
+        }
+    }
+}
+
+/// Returns only the business dates (Mon-Fri) between start_date and end_date
+/// that match the desired frequency.
+pub fn get_bdates_list_with_freq(
+    start_date_str: &str,
+    end_date_str: &str,
+    freq: BDateFreq,
+) -> Result<Vec<NaiveDate>, Box<dyn Error>> {
+    let start_date = NaiveDate::parse_from_str(start_date_str, "%Y-%m-%d")?;
+    let end_date = NaiveDate::parse_from_str(end_date_str, "%Y-%m-%d")?;
+
+    if start_date > end_date {
+        return Ok(Vec::new());
+    }
+
+    let mut dates = match freq {
+        BDateFreq::Daily => collect_daily(start_date, end_date),
+
+        BDateFreq::WeeklyMonday => collect_weekly(start_date, end_date, Weekday::Mon),
+        BDateFreq::WeeklyFriday => collect_weekly(start_date, end_date, Weekday::Fri),
+
+        BDateFreq::MonthStart => collect_monthly(start_date, end_date, /*start=*/ true),
+        BDateFreq::MonthEnd => collect_monthly(start_date, end_date, /*start=*/ false),
+
+        BDateFreq::QuarterStart => collect_quarterly(start_date, end_date, /*start=*/ true),
+        BDateFreq::QuarterEnd => collect_quarterly(start_date, end_date, /*start=*/ false),
+
+        BDateFreq::YearStart => collect_yearly(start_date, end_date, /*start=*/ true),
+        BDateFreq::YearEnd => collect_yearly(start_date, end_date, /*start=*/ false),
+    };
+
+    // Filter out any weekend days that might slip in edge cases (e.g. if the
+    // computed "start of month" fell on Sat/Sun).
+    dates.retain(|d| d.weekday() != Weekday::Sat && d.weekday() != Weekday::Sun);
+
+    Ok(dates)
+}
+
+/* ------------------------------ Helpers ------------------------------ */
+
+/// Return all business days, day-by-day.
+fn collect_daily(start_date: NaiveDate, end_date: NaiveDate) -> Vec<NaiveDate> {
+    let mut result = Vec::new();
+    let mut current = start_date;
+    while current <= end_date {
+        if is_weekday(current) {
+            result.push(current);
+        }
+        current = current.succ_opt().unwrap();
+    }
+    result
+}
+
+/// Return the specified weekday (e.g. Monday, Friday) in each week of the range.
+fn collect_weekly(
+    start_date: NaiveDate,
+    end_date: NaiveDate,
+    target_weekday: Weekday,
+) -> Vec<NaiveDate> {
+    let mut result = Vec::new();
+
+    // Find the first `target_weekday` on or after `start_date`.
+    // If `start_date` is already e.g. Monday, we can use it as is.
+    // Otherwise, jump ahead until we get that weekday.
+    let mut current = move_to_weekday_on_or_after(start_date, target_weekday);
+
+    // Step in 7-day increments (full weeks).
+    while current <= end_date {
+        result.push(current);
+        current = current + Duration::days(7);
+    }
+    result
+}
+
+/// Return either first or last business day in each month of the range.
+fn collect_monthly(
+    start_date: NaiveDate,
+    end_date: NaiveDate,
+    want_first_day: bool,
+) -> Vec<NaiveDate> {
+    let mut result = Vec::new();
+
+    // We'll iterate month by month, from (start_year, start_month) up to
+    // (end_year, end_month).
+    let mut year = start_date.year();
+    let mut month = start_date.month();
+
+    // A small helper that updates year/month by +1 month.
+    let next_month = |(yr, mo): (i32, u32)| -> (i32, u32) {
+        if mo == 12 {
+            (yr + 1, 1)
+        } else {
+            (yr, mo + 1)
+        }
+    };
+
+    // Move `(year, month)` backward if necessary so that `(year, month)`
+    // definitely covers the entire period from `start_date` onward.
+    // Actually, it’s simpler to start from the actual (year, month) of start_date
+    // and go up. We'll just skip if the computed "day" < start_date.
+
+    // Continue while we haven't passed (end_year, end_month).
+    while year < end_date.year() || (year == end_date.year() && month <= end_date.month()) {
+        // Compute the date that represents either first or last business day
+        // for this (year, month).
+        let candidate = if want_first_day {
+            first_business_day_of_month(year, month)
+        } else {
+            last_business_day_of_month(year, month)
+        };
+        if candidate >= start_date && candidate <= end_date {
+            result.push(candidate);
+        }
+
+        // Move to the next month.
+        let (ny, nm) = next_month((year, month));
+        year = ny;
+        month = nm;
+    }
+
+    result
+}
+
+/// Return either the first or last business day in each quarter of the range.
+fn collect_quarterly(
+    start_date: NaiveDate,
+    end_date: NaiveDate,
+    want_first_day: bool,
+) -> Vec<NaiveDate> {
+    let mut result = Vec::new();
+
+    // We'll figure out which quarter `start_date` is in, then jump quarter-by-quarter.
+    // Quarters are: Q1 = months 1–3, Q2 = 4–6, Q3 = 7–9, Q4 = 10–12.
+    // Start by computing the (year, quarter_index) for start_date.
+    let mut year = start_date.year();
+    let mut q = month_to_quarter(start_date.month());
+
+    while quarter_to_first_date(year, q) > end_date {
+        // If even the earliest day in that quarter is > end_date, we’re done.
+        return result;
+    }
+
+    // Move backward if the quarter’s last day < start_date, etc.
+    // But simpler: we’ll do a loop that increments quarter by quarter, and
+    // pick the appropriate date each time. We break when we pass end_date.
+
+    loop {
+        // For the current year+quarter, compute the date that’s either the first or last
+        // business day of that quarter:
+        let candidate = if want_first_day {
+            first_business_day_of_quarter(year, q)
+        } else {
+            last_business_day_of_quarter(year, q)
+        };
+
+        if candidate > end_date {
+            break;
+        }
+        if candidate >= start_date {
+            result.push(candidate);
+        }
+
+        // Move to next quarter.
+        if q == 4 {
+            year += 1;
+            q = 1;
+        } else {
+            q += 1;
+        }
+    }
+
+    result
+}
+
+/// Return either the first or last business day in each year of the range.
+fn collect_yearly(
+    start_date: NaiveDate,
+    end_date: NaiveDate,
+    want_first_day: bool,
+) -> Vec<NaiveDate> {
+    let mut result = Vec::new();
+    let mut year = start_date.year();
+
+    // Step year-by-year from `start_date.year()` up to `end_date.year()`.
+    while year <= end_date.year() {
+        let candidate = if want_first_day {
+            first_business_day_of_year(year)
+        } else {
+            last_business_day_of_year(year)
+        };
+        if candidate >= start_date && candidate <= end_date {
+            result.push(candidate);
+        }
+        year += 1;
+    }
+    result
+}
+
+/* ---------------------- Low-Level Utility Functions ---------------------- */
+
+/// Is this a weekday (Mon-Fri)?
+fn is_weekday(date: NaiveDate) -> bool {
+    match date.weekday() {
+        Weekday::Sat | Weekday::Sun => false,
+        _ => true,
+    }
+}
+
+/// 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 mut current = date;
+    while current.weekday() != target {
+        current = current.succ_opt().unwrap();
+    }
+    current
+}
+
+/// Return the earliest business day of (year, month).
+fn first_business_day_of_month(year: i32, month: u32) -> NaiveDate {
+    // Start with the 1st of the month.
+    let mut d = NaiveDate::from_ymd_opt(year, month, 1).expect("invalid year-month");
+    // If it’s Sat/Sun, move forward until we get a weekday.
+    while !is_weekday(d) {
+        d = d.succ_opt().unwrap();
+    }
+    d
+}
+
+/// Return the latest business day of (year, month).
+fn last_business_day_of_month(year: i32, month: u32) -> NaiveDate {
+    let last_dom = days_in_month(year, month);
+    let mut d = NaiveDate::from_ymd_opt(year, month, last_dom).expect("invalid year-month");
+    // If it’s Sat/Sun, move backward until we get a weekday.
+    while !is_weekday(d) {
+        d = d.pred_opt().unwrap();
+    }
+    d
+}
+
+/// Number of days in a month (not considering leap years *beyond* chrono's normal handling).
+fn days_in_month(year: i32, month: u32) -> u32 {
+    // Chrono can handle this if we do a little trick:
+    // Construct the 1st of the next month, then subtract 1 day.
+    // For example:
+    // if month == 12 => next = (year+1, 1, 1)
+    // else => next = (year, month+1, 1)
+    let (ny, nm) = if month == 12 {
+        (year + 1, 1)
+    } else {
+        (year, month + 1)
+    };
+    let first_of_next = NaiveDate::from_ymd_opt(ny, nm, 1).unwrap();
+    let last_of_this = first_of_next.pred_opt().unwrap();
+    last_of_this.day()
+}
+
+/// Convert a month (1..12) to a quarter (1..4).
+fn month_to_quarter(m: u32) -> u32 {
+    (m - 1) / 3 + 1
+}
+
+/// Returns 1st day of a given (year, quarter).
+fn quarter_to_first_date(year: i32, quarter: u32) -> NaiveDate {
+    let month = match quarter {
+        1 => 1,
+        2 => 4,
+        3 => 7,
+        4 => 10,
+        _ => panic!("invalid quarter"),
+    };
+    NaiveDate::from_ymd_opt(year, month, 1).unwrap()
+}
+
+/// Return the earliest business day in (year, quarter).
+fn first_business_day_of_quarter(year: i32, quarter: u32) -> NaiveDate {
+    let mut d = quarter_to_first_date(year, quarter);
+    while !is_weekday(d) {
+        d = d.succ_opt().unwrap();
+    }
+    d
+}
+
+/// Return the last business day in (year, quarter).
+fn last_business_day_of_quarter(year: i32, quarter: u32) -> NaiveDate {
+    // The last month in the quarter is quarter_to_first_date(...) + 2 months
+    // Then we find the last day of that month.
+    let start = quarter_to_first_date(year, quarter);
+    let last_month = start.month() + 2; // e.g. Q1 => month=1 => +2=3 => March
+    last_business_day_of_month(year, last_month)
+}
+
+/// Returns Jan 1st of a given year (adjust if weekend).
+fn first_business_day_of_year(year: i32) -> NaiveDate {
+    let mut d = NaiveDate::from_ymd_opt(year, 1, 1).unwrap();
+    while !is_weekday(d) {
+        d = d.succ_opt().unwrap();
+    }
+    d
+}
+
+/// Returns Dec 31st of a given year (adjust if weekend).
+fn last_business_day_of_year(year: i32) -> NaiveDate {
+    let mut d = NaiveDate::from_ymd_opt(year, 12, 31).unwrap();
+    while !is_weekday(d) {
+        d = d.pred_opt().unwrap();
+    }
+    d
+}