From 93f88ab537d29d038d5ccf78799fe4e7b62dab8c Mon Sep 17 00:00:00 2001 From: Palash Tyagi <23239946+Magnus167@users.noreply.github.com> Date: Sun, 13 Apr 2025 11:10:07 +0100 Subject: [PATCH] Refactor date utilities: remove BDateFreq and AggregationType enums, integrate frequency handling in get_bdates_series_default_opt --- src/utils/dateutils.rs | 82 +++++++++++------------------------------- 1 file changed, 21 insertions(+), 61 deletions(-) diff --git a/src/utils/dateutils.rs b/src/utils/dateutils.rs index 2261c98..095fce6 100644 --- a/src/utils/dateutils.rs +++ b/src/utils/dateutils.rs @@ -1,3 +1,5 @@ +use crate::utils::bdates; +use crate::utils::bdates::BDateFreq; use chrono::NaiveDate; use chrono::{Datelike, Weekday}; use polars::prelude::*; @@ -57,61 +59,6 @@ pub fn get_bdates_list( Ok(business_days) } -#[derive(Debug, Clone, Copy)] -pub enum BDateFreq { - Daily, - WeeklyMonday, - MonthStart, - QuarterStart, - YearStart, - MonthEnd, - QuarterEnd, - WeeklyFriday, - YearEnd, -} - -impl BDateFreq { - pub fn from_string(freq: String) -> Result> { - // use `from_str` to convert the string to a BDateFreq enum - Self::from_str(&freq) - } - pub fn from_str(freq: &str) -> Result> { - 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, - } - } -} - -#[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. -} - -// Map a BDateFreq to an AggregationType. fn compute_group_key(d: NaiveDate, freq: BDateFreq) -> String { match freq { // For Daily, each date is its own group. @@ -134,19 +81,32 @@ fn compute_group_key(d: NaiveDate, freq: BDateFreq) -> String { BDateFreq::YearStart | BDateFreq::YearEnd => format!("{}", d.year()), } } - -pub fn get_bdates_series_default( +pub fn get_bdates_series_default_opt( start_date: String, end_date: String, freq: Option, ) -> Result> { let freq = freq.unwrap_or_else(|| "D".to_string()); let freq = BDateFreq::from_str(&freq)?; - get_bdates_series(start_date, end_date, freq) + let series = Series::new( + "bdates".into(), + bdates::get_bdates_list_with_freq(&start_date, &end_date, freq)?, + ); + Ok(series) +} + +pub fn get_bdates_series_default_pl( + start_date: String, + end_date: String, + freq: Option, +) -> Result> { + let freq = freq.unwrap_or_else(|| "D".to_string()); + let freq = BDateFreq::from_str(&freq)?; + get_bdates_series_pl(start_date, end_date, freq) } /// Get the business dates between two dates as a Series. -pub fn get_bdates_series( +pub fn get_bdates_series_pl( start_date: String, end_date: String, freq: BDateFreq, @@ -163,8 +123,8 @@ pub fn get_bdates_series( ])?; let gb = df.lazy().group_by(["group"]); let aggx = match freq.agg_type() { - AggregationType::Start => gb.agg([col("bdates").first()]), - AggregationType::End => gb.agg([col("bdates").last()]), + bdates::AggregationType::Start => gb.agg([col("bdates").first()]), + bdates::AggregationType::End => gb.agg([col("bdates").last()]), }; let result = aggx.collect()?; let result = result