feat: ticker format, test_utils, Makefile

CONTRIBUTING.md

@@ -1,6 +1,5 @@
 # Contributing code
-- Make sure the test pass
-- Run `cargo clippy --fix --allow-dirty`
+- Run `make precommit`
 
 # Dev Install
 ## Dev Containers

Cargo.toml

@@ -15,7 +15,7 @@ polars = { version = "0.32", features = [
 	"object",
 	"dtype-struct",
 ] }
-reqwest = { version = "0.11", features = ["blocking", "gzip"] }
+reqwest = { version = "0.12", features = ["blocking", "gzip"] }
 glob = { version = "0.3" }
 clokwerk = "0.4"
 strum_macros = "0.26"
@@ -30,6 +30,6 @@ lazy_static = "1.4"
 anyhow = "1.0"
 
 [dev-dependencies]
-serial_test = "*"
-rstest = "0.21"
+serial_test = "3.1"
+rstest = "0.23"
 pretty_assertions = "1.4"

Makefile

@@ -0,0 +1,7 @@
+precommit:
+	rustup update
+	cargo update
+	cargo check
+	cargo fmt
+	cargo t
+	cargo clippy --fix --allow-dirty

src/lib.rs

@@ -1,2 +1,5 @@
 pub mod util;
 pub use util::*;
+
+#[cfg(test)]
+mod test_utils;

src/main.rs

@@ -1,4 +1,5 @@
 use anyhow::{Error, Result};
+use ark_invest_api_rust_data::{util::ticker::Ticker, *};
 use clokwerk::{AsyncScheduler, Job, TimeUnits};
 use futures::future::join_all;
 use lazy_static::lazy_static;
@@ -11,10 +12,6 @@ use strum::IntoEnumIterator;
 use tokio::task;
 use tokio::time::Duration;
 
-mod util;
-use util::ticker::Ticker;
-use util::*;
-
 lazy_static! {
     static ref SOURCE: Source = match env::var("ARK_SOURCE") {
         Ok(val) =>

src/test_utils.rs

@@ -0,0 +1,17 @@
+use anyhow::{Error, Result};
+use polars::prelude::*;
+
+pub fn defualt_df(ticker: &[Option<&str>], company: &[Option<&str>]) -> Result<DataFrame, Error> {
+    let target_len = ticker.len() + 1;
+    let df = df![
+        "date" => vec!["2024-01-01"; target_len],
+        "ticker" => [ticker, &[Some("TSLA")]].concat(),
+        "cusip" => vec!["TESLA"; target_len],
+        "company" => [company, &[Some("TESLA")]].concat(),
+        "market_value" => vec![10; target_len],
+        "shares" => vec![10; target_len],
+        "share_price" => vec![100.00; target_len],
+        "weight" => vec![10.00; target_len],
+    ]?;
+    Ok(df)
+}

src/util.rs

@@ -13,7 +13,7 @@ use strum_macros::EnumString;
 use ticker::{DataSource, Ticker};
 pub mod data_reader;
 pub mod df;
-mod df_format;
+mod format;
 pub mod ticker;
 
 #[derive(Debug, Default, EnumString, Clone, Copy, PartialEq)]
@@ -159,14 +159,14 @@ impl Ark {
         let mut df = df.collect()?;
         match data_source {
             Some(ds) => {
-                df = df_format::df_format(ds, df.into())?.collect()?;
+                df = format::data_source(ds, df.into())?.collect()?;
             }
             None => {
-                df = df_format::df_format_europe_csv(df.into())?.collect()?;
-                df = df_format::df_format_europe_arkfundsio(df.into())?.collect()?;
-                df = df_format::df_format_21shares(df.into())?.collect()?;
-                df = df_format::df_format_arkvx(df.into())?.collect()?;
-                df = df_format::df_format_europe(df.into())?.collect()?;
+                df = format::df_format_europe_csv(df.into())?.collect()?;
+                df = format::df_format_europe_arkfundsio(df.into())?.collect()?;
+                df = format::df_format_21shares(df.into())?.collect()?;
+                df = format::df_format_arkvx(df.into())?.collect()?;
+                df = format::df_format_europe(df.into())?.collect()?;
             }
         }
 
@@ -242,34 +242,7 @@ impl Ark {
             }
         }
 
-        // format arkw, ARK BITCOIN ETF HOLDCO (ARKW) to ARKB
-        if let Ok(x) = df
-            .clone()
-            .lazy()
-            .with_columns(vec![
-                when(col("company").eq(lit("ARK BITCOIN ETF HOLDCO (ARKW)")))
-                    .then(lit("ARKB"))
-                    .otherwise(col("ticker"))
-                    .alias("ticker"),
-                when(col("company").eq(lit("ARK BITCOIN ETF HOLDCO (ARKW)")))
-                    .then(lit("ARKB"))
-                    .otherwise(col("company"))
-                    .alias("company"),
-            ])
-            .with_columns(vec![
-                when(col("company").eq(lit("ARK BITCOIN ETF HOLDCO (ARKF)")))
-                    .then(lit("ARKB"))
-                    .otherwise(col("ticker"))
-                    .alias("ticker"),
-                when(col("company").eq(lit("ARK BITCOIN ETF HOLDCO (ARKF)")))
-                    .then(lit("ARKB"))
-                    .otherwise(col("company"))
-                    .alias("company"),
-            ])
-            .collect()
-        {
-            df = x;
-        }
+        df = format::Ticker::all(df.into())?.collect()?;
 
         let mut expressions: Vec<Expr> = vec![];
 
@@ -581,23 +554,15 @@ impl Ark {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::test_utils::*;
     use pretty_assertions::assert_eq;
     use serial_test::serial;
     use std::fs;
 
     #[test]
     #[serial]
     fn read_write_parquet() -> Result<(), Error> {
-        let test_df = df![
-            "date" => ["2023-01-01"],
-            "ticker" => ["TSLA"],
-            "cusip" => ["123abc"],
-            "company" => ["Tesla"],
-            "market_value" => [100],
-            "shares" => [10],
-            "share_price" => [10],
-            "weight" => [10.00]
-        ]?;
+        let test_df = defualt_df(&[Some("COIN")], &[Some("COINBASE")])?;
 
         Ark::write_df_parquet("data/test/ARKK.parquet".into(), test_df.clone().into())?;
         let read = Ark::new(Source::Read, Ticker::ARKK, Some("data/test".to_owned()))?.collect()?;
@@ -610,16 +575,14 @@ mod tests {
     #[test]
     #[serial]
     fn arkw_format_arkb() -> Result<(), Error> {
-        let test_df = df![
-            "date" => ["2024-01-01", "2024-01-02"],
-            "ticker" => [None::<&str>, Some("TSLA")],
-            "cusip" => ["123abc", "TESLA"],
-            "company" => ["ARK BITCOIN ETF HOLDCO (ARKW)", "TESLA"],
-            "market_value" => [100, 400],
-            "shares" => [10, 20],
-            "share_price" => [10, 20],
-            "weight" => [10.00, 20.00]
-        ]?;
+        let test_df = defualt_df(
+            &[None::<&str>, Some("ARKB"), Some("ARKB")],
+            &[
+                Some("ARK BITCOIN ETF HOLDCO (ARKW)"),
+                Some("ARK BITCOIN ETF HOLDCO (ARKW)"),
+                Some("ARKB"),
+            ],
+        )?;
 
         Ark::write_df_parquet("data/test/ARKW.parquet".into(), test_df.clone().into())?;
         let read = Ark::new(Source::Read, Ticker::ARKW, Some("data/test".to_owned()))?.collect()?;
@@ -628,16 +591,10 @@ mod tests {
         let df = Ark::df_format(read.into(), None)?.collect()?;
         assert_eq!(
             df,
-            df![
-                "date" => ["2024-01-01", "2024-01-02"],
-                "ticker" => ["ARKB", "TSLA"],
-                "cusip" => ["123abc", "TESLA"],
-                "company" => ["ARKB", "TESLA"],
-                "market_value" => [100, 400],
-                "shares" => [10, 20],
-                "share_price" => [10, 20],
-                "weight" => [10.00, 20.00]
-            ]?
+            defualt_df(
+                &[Some("ARKB"), Some("ARKB"), Some("ARKB")],
+                &[Some("ARKB"), Some("ARKB"), Some("ARKB")]
+            )?,
         );
 
         Ok(())
@@ -646,34 +603,25 @@ mod tests {
     #[test]
     #[serial]
     fn arkf_format_arkb() -> Result<(), Error> {
-        let test_df = df![
-            "date" => ["2024-01-01", "2024-01-02"],
-            "ticker" => [None::<&str>, Some("TSLA")],
-            "cusip" => ["123abc", "TESLA"],
-            "company" => ["ARK BITCOIN ETF HOLDCO (ARKF)", "TESLA"],
-            "market_value" => [100, 400],
-            "shares" => [10, 20],
-            "share_price" => [10, 20],
-            "weight" => [10.00, 20.00]
-        ]?;
-
+        let test_df = defualt_df(
+            &[None::<&str>, Some("ARKB"), Some("ARKB")],
+            &[
+                Some("ARK BITCOIN ETF HOLDCO (ARKF)"),
+                Some("ARK BITCOIN ETF HOLDCO (ARKF)"),
+                Some("ARKB"),
+            ],
+        )?;
         Ark::write_df_parquet("data/test/ARKF.parquet".into(), test_df.clone().into())?;
         let read = Ark::new(Source::Read, Ticker::ARKF, Some("data/test".to_owned()))?.collect()?;
         fs::remove_file("data/test/ARKF.parquet")?;
 
         let df = Ark::df_format(read.into(), None)?.collect()?;
         assert_eq!(
             df,
-            df![
-                "date" => ["2024-01-01", "2024-01-02"],
-                "ticker" => ["ARKB", "TSLA"],
-                "cusip" => ["123abc", "TESLA"],
-                "company" => ["ARKB", "TESLA"],
-                "market_value" => [100, 400],
-                "shares" => [10, 20],
-                "share_price" => [10, 20],
-                "weight" => [10.00, 20.00]
-            ]?
+            defualt_df(
+                &[Some("ARKB"), Some("ARKB"), Some("ARKB")],
+                &[Some("ARKB"), Some("ARKB"), Some("ARKB")]
+            )?,
         );
 
         Ok(())

src/util/df_format.rs → src/util/format/data_source.rs

@@ -3,7 +3,7 @@ use polars::prelude::*;
 
 use crate::{ticker::DataSource, util::df::DF};
 
-pub fn df_format(data_source: DataSource, mut df: DF) -> Result<DF, Error> {
+pub fn data_source(data_source: DataSource, mut df: DF) -> Result<DF, Error> {
     let df = match data_source {
         DataSource::ArkVenture => df_format_arkvx(df)?,
         DataSource::Ark => df,

src/util/format/mod.rs

@@ -0,0 +1,4 @@
+pub mod data_source;
+pub use data_source::*;
+pub mod ticker;
+pub use ticker::*;