diff --git a/src/util/format/ticker.rs b/src/util/format/ticker.rs index 6bcb8b6..63ccac5 100644 --- a/src/util/format/ticker.rs +++ b/src/util/format/ticker.rs @@ -17,6 +17,7 @@ pub enum Ticker { XYZ, CASH_USD, TSM, + RKLB, } impl Ticker { @@ -38,6 +39,7 @@ impl Ticker { Self::XYZ => Self::xyz(df), Self::CASH_USD => Self::cash_usd(df), Self::TSM => Self::tsm(df), + Self::RKLB => Self::rklb(df), } } @@ -250,6 +252,30 @@ impl Ticker { Ok(df.into()) } + + fn rklb(df: DF) -> Result { + let mut df = df.collect()?; + + if let Ok(x) = df + .clone() + .lazy() + .with_columns(vec![ + when(col("company").eq(lit("ROCKET LAB"))) + .then(lit("RKLB")) + .otherwise(col("ticker")) + .alias("ticker"), + when(col("company").eq(lit("ROCKET LAB USA"))) + .then(lit("ROCKET LAB")) + .otherwise(col("company")) + .alias("company") + ]) + .collect() + { + df = x; + } + + Ok(df.into()) + } } #[cfg(test)] @@ -364,6 +390,17 @@ mod tests { &[Some("TMSC")] )?, )] + #[case::rklb( + Ticker::RKLB, + defualt_df( + &[Some("RKLB"), Some("RKLB"), None::<&str>], + &[Some("ROCKET LAB"), Some("ROCKET LAB USA"), Some("ROCKET LAB")], + )?, + defualt_df( + &[Some("RKLB"), Some("RKLB"), Some("RKLB")], + &[Some("ROCKET LAB"), Some("ROCKET LAB"), Some("ROCKET LAB")] + )?, + )] fn matrix( #[case] ticker: Ticker, #[case] input: DataFrame,