Refactor date utilities: remove BDateFreq and AggregationType enums, integrate frequency handling in get_bdates_series_default_opt

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<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,
-        }
-    }
-}
-
-#[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<String>,
 ) -> Result<Series, Box<dyn Error>> {
     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<String>,
+) -> Result<Series, Box<dyn Error>> {
+    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