use std::sync::Arc;
use polars::prelude::{DataType, Field, PlSmallStr, Schema, SchemaRef, TimeUnit, TimeZone};
use strum::{Display, EnumString, IntoStaticStr};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Display, EnumString, IntoStaticStr)]
#[strum(serialize_all = "snake_case")]
pub enum CanonicalCol {
DataSource,
CountryCode,
CurrencyCode,
Timestamp,
OpenTimestamp,
Period,
Open,
High,
Low,
Close,
Price,
Volume,
TradeId,
QuoteAssetVolume,
NumberOfTrades,
TakerBuyBaseAssetVolume,
TakerBuyQuoteAssetVolume,
TakerSellBaseAssetVolume,
TakerSellQuoteAssetVolume,
NumberOfBuyTrades,
NumberOfSellTrades,
IsBuyerMaker,
IsBestMatch,
PriceBinStart,
PriceBinEnd,
TimeSlotCount,
Category,
NewsType,
NewsTypeConfidence,
NewsTypeSource,
NewsName,
EconomicImpact,
Actual,
Forecast,
Previous,
}
impl From<CanonicalCol> for PlSmallStr {
fn from(value: CanonicalCol) -> Self {
value.as_str().into()
}
}
impl CanonicalCol {
pub fn name(&self) -> PlSmallStr {
(*self).into()
}
pub fn as_str(&self) -> &'static str {
self.into()
}
pub fn dtype(&self) -> DataType {
match self {
Self::DataSource
| Self::CountryCode
| Self::CurrencyCode
| Self::Category
| Self::NewsType
| Self::NewsTypeSource
| Self::NewsName
| Self::Period => DataType::String,
Self::TradeId
| Self::NumberOfTrades
| Self::EconomicImpact
| Self::NumberOfBuyTrades
| Self::NumberOfSellTrades
| Self::TimeSlotCount => DataType::Int64,
Self::Open
| Self::High
| Self::Low
| Self::Close
| Self::Price
| Self::Volume
| Self::QuoteAssetVolume
| Self::TakerBuyBaseAssetVolume
| Self::TakerBuyQuoteAssetVolume
| Self::TakerSellBaseAssetVolume
| Self::TakerSellQuoteAssetVolume
| Self::PriceBinStart
| Self::PriceBinEnd
| Self::Actual
| Self::Forecast
| Self::Previous
| Self::NewsTypeConfidence => DataType::Float64,
Self::Timestamp | Self::OpenTimestamp => {
DataType::Datetime(TimeUnit::Microseconds, Some(TimeZone::UTC))
}
Self::IsBuyerMaker | Self::IsBestMatch => DataType::Boolean,
}
}
pub fn field(&self) -> Field {
Field::new(self.name(), self.dtype())
}
}
pub fn ohlcv_future_schema() -> SchemaRef {
let s = Schema::from_iter([
CanonicalCol::OpenTimestamp.field(),
CanonicalCol::Open.field(),
CanonicalCol::High.field(),
CanonicalCol::Low.field(),
CanonicalCol::Close.field(),
CanonicalCol::Volume.field(),
CanonicalCol::Timestamp.field(), ]);
Arc::new(s)
}
pub fn ohlcv_spot_schema() -> SchemaRef {
let s = Schema::from_iter([
CanonicalCol::OpenTimestamp.field(),
CanonicalCol::Open.field(),
CanonicalCol::High.field(),
CanonicalCol::Low.field(),
CanonicalCol::Close.field(),
CanonicalCol::Volume.field(),
CanonicalCol::Timestamp.field(), CanonicalCol::QuoteAssetVolume.field(),
CanonicalCol::NumberOfTrades.field(),
CanonicalCol::TakerBuyBaseAssetVolume.field(),
CanonicalCol::TakerBuyQuoteAssetVolume.field(),
]);
Arc::new(s)
}
pub fn trades_spot_schema() -> SchemaRef {
let s = Schema::from_iter([
CanonicalCol::TradeId.field(),
CanonicalCol::Price.field(),
CanonicalCol::Volume.field(), CanonicalCol::QuoteAssetVolume.field(), CanonicalCol::Timestamp.field(), CanonicalCol::IsBuyerMaker.field(),
CanonicalCol::IsBestMatch.field(),
]);
Arc::new(s)
}
pub fn volume_profile_spot_schema() -> SchemaRef {
let s = Schema::from_iter([
CanonicalCol::OpenTimestamp.field(), CanonicalCol::Timestamp.field(), CanonicalCol::PriceBinStart.field(),
CanonicalCol::PriceBinEnd.field(),
CanonicalCol::Volume.field(), CanonicalCol::TakerBuyBaseAssetVolume.field(),
CanonicalCol::TakerSellBaseAssetVolume.field(),
CanonicalCol::QuoteAssetVolume.field(), CanonicalCol::TakerBuyQuoteAssetVolume.field(),
CanonicalCol::TakerSellQuoteAssetVolume.field(),
CanonicalCol::NumberOfTrades.field(),
CanonicalCol::NumberOfBuyTrades.field(),
CanonicalCol::NumberOfSellTrades.field(),
]);
Arc::new(s)
}
pub fn tpo_spot_schema() -> SchemaRef {
let s = Schema::from_iter([
CanonicalCol::OpenTimestamp.field(), CanonicalCol::Timestamp.field(), CanonicalCol::PriceBinStart.field(),
CanonicalCol::PriceBinEnd.field(),
CanonicalCol::TimeSlotCount.field(),
]);
Arc::new(s)
}
pub fn tpo_future_schema() -> SchemaRef {
let s = Schema::from_iter([
CanonicalCol::OpenTimestamp.field(), CanonicalCol::Timestamp.field(), CanonicalCol::PriceBinStart.field(),
CanonicalCol::PriceBinEnd.field(),
CanonicalCol::TimeSlotCount.field(),
]);
Arc::new(s)
}
pub fn economic_calendar_schema() -> SchemaRef {
let s = Schema::from_iter([
CanonicalCol::DataSource.field(),
CanonicalCol::Category.field(),
CanonicalCol::Timestamp.field(),
CanonicalCol::NewsType.field(),
CanonicalCol::NewsTypeConfidence.field(), CanonicalCol::NewsTypeSource.field(), CanonicalCol::Period.field(), CanonicalCol::NewsName.field(),
CanonicalCol::CountryCode.field(),
CanonicalCol::CurrencyCode.field(),
CanonicalCol::EconomicImpact.field(),
CanonicalCol::Actual.field(),
CanonicalCol::Forecast.field(),
CanonicalCol::Previous.field(),
]);
Arc::new(s)
}