use std::sync::Arc;
use polars::prelude::{DataType, Field, PlSmallStr, Schema, SchemaRef, TimeUnit, TimeZone};
use strum::{AsRefStr, Display, EnumString, IntoStaticStr};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Display, EnumString, IntoStaticStr, AsRefStr)]
#[strum(serialize_all = "snake_case")]
pub enum CanonicalCol {
DataSource,
CountryCode,
CurrencyCode,
Date,
PointInTime,
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,
Roc,
RocAbsolute,
SessionDate,
SessionHigh,
SessionLow,
SessionHighestClose,
SessionLowestClose,
SessionVolume,
SessionVwap,
}
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 const 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
| Self::Roc
| Self::RocAbsolute
| Self::SessionHigh
| Self::SessionLow
| Self::SessionHighestClose
| Self::SessionLowestClose
| Self::SessionVolume
| Self::SessionVwap => DataType::Float64,
Self::Date | Self::SessionDate => DataType::Date,
Self::PointInTime | 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::PointInTime.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::PointInTime.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::PointInTime.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::PointInTime.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::PointInTime.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::PointInTime.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::PointInTime.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)
}