boris
97 lines
2.5 KiB
Rust
97 lines
2.5 KiB
Rust
use std::collections::HashMap;
|
|
|
|
use chrono::NaiveDate;
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct TradingCalendar {
|
|
days: Vec<NaiveDate>,
|
|
index: HashMap<NaiveDate, usize>,
|
|
}
|
|
|
|
impl TradingCalendar {
|
|
pub fn new(mut days: Vec<NaiveDate>) -> Self {
|
|
days.sort_unstable();
|
|
days.dedup();
|
|
|
|
let index = days
|
|
.iter()
|
|
.copied()
|
|
.enumerate()
|
|
.map(|(idx, day)| (day, idx))
|
|
.collect();
|
|
|
|
Self { days, index }
|
|
}
|
|
|
|
pub fn days(&self) -> &[NaiveDate] {
|
|
&self.days
|
|
}
|
|
|
|
pub fn iter(&self) -> impl Iterator<Item = NaiveDate> + '_ {
|
|
self.days.iter().copied()
|
|
}
|
|
|
|
pub fn len(&self) -> usize {
|
|
self.days.len()
|
|
}
|
|
|
|
pub fn is_empty(&self) -> bool {
|
|
self.days.is_empty()
|
|
}
|
|
|
|
pub fn index_of(&self, date: NaiveDate) -> Option<usize> {
|
|
self.index.get(&date).copied()
|
|
}
|
|
|
|
pub fn previous_day(&self, date: NaiveDate) -> Option<NaiveDate> {
|
|
let idx = self.index_of(date)?;
|
|
idx.checked_sub(1)
|
|
.and_then(|prev| self.days.get(prev).copied())
|
|
}
|
|
|
|
pub fn previous_trading_date(&self, date: NaiveDate, n: usize) -> Option<NaiveDate> {
|
|
if n == 0 {
|
|
return None;
|
|
}
|
|
let before_count = match self.days.binary_search(&date) {
|
|
Ok(idx) => idx,
|
|
Err(idx) => idx,
|
|
};
|
|
before_count
|
|
.checked_sub(n)
|
|
.and_then(|idx| self.days.get(idx).copied())
|
|
}
|
|
|
|
pub fn next_trading_date(&self, date: NaiveDate, n: usize) -> Option<NaiveDate> {
|
|
if n == 0 {
|
|
return None;
|
|
}
|
|
let first_after = match self.days.binary_search(&date) {
|
|
Ok(idx) => idx.saturating_add(1),
|
|
Err(idx) => idx,
|
|
};
|
|
first_after
|
|
.checked_add(n.saturating_sub(1))
|
|
.and_then(|idx| self.days.get(idx).copied())
|
|
}
|
|
|
|
pub fn trading_dates(&self, start: NaiveDate, end: NaiveDate) -> Vec<NaiveDate> {
|
|
if start > end {
|
|
return Vec::new();
|
|
}
|
|
self.days
|
|
.iter()
|
|
.copied()
|
|
.filter(|date| *date >= start && *date <= end)
|
|
.collect()
|
|
}
|
|
|
|
pub fn trailing_days(&self, end: NaiveDate, lookback: usize) -> Vec<NaiveDate> {
|
|
let Some(end_idx) = self.index_of(end) else {
|
|
return Vec::new();
|
|
};
|
|
let start = end_idx.saturating_add(1).saturating_sub(lookback);
|
|
self.days[start..=end_idx].to_vec()
|
|
}
|
|
}
|