brk_query 0.3.0-beta.9

An interface to find and format data from BRK
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

use std::{fs, path::PathBuf};

use brk_error::{Error, Result};
use brk_types::{Cohort, Date, Day1, Urpd, UrpdAggregation, UrpdRaw};
use vecdb::ReadableOptionVec;

use crate::Query;

impl Query {
    /// Available cohorts for URPD.
    pub fn urpd_cohorts(&self) -> Result<Vec<Cohort>> {
        let states_path = &self.computer().distribution.states_path;

        let mut cohorts: Vec<Cohort> = fs::read_dir(states_path)?
            .filter_map(|entry| {
                let name = entry.ok()?.file_name().into_string().ok()?;
                if !states_path.join(&name).join("urpd").exists() {
                    return None;
                }
                Cohort::new(name)
            })
            .collect();

        cohorts.sort_unstable();

        Ok(cohorts)
    }

    pub(crate) fn urpd_dir(&self, cohort: &Cohort) -> Result<PathBuf> {
        let dir = self
            .computer()
            .distribution
            .states_path
            .join(cohort)
            .join("urpd");

        if !dir.exists() {
            let valid = self
                .urpd_cohorts()
                .unwrap_or_default()
                .iter()
                .map(|c| c.to_string())
                .collect::<Vec<_>>()
                .join(", ");
            return Err(Error::NotFound(format!(
                "Unknown cohort '{cohort}'. Available: {valid}"
            )));
        }

        Ok(dir)
    }

    /// Available dates for a cohort.
    pub fn urpd_dates(&self, cohort: &Cohort) -> Result<Vec<Date>> {
        let dir = self.urpd_dir(cohort)?;

        let mut dates: Vec<Date> = fs::read_dir(&dir)?
            .filter_map(|entry| entry.ok()?.file_name().to_str()?.parse().ok())
            .collect();

        dates.sort_unstable();
        Ok(dates)
    }

    /// Raw URPD data for a cohort on a specific date.
    pub fn urpd_raw(&self, cohort: &Cohort, date: Date) -> Result<UrpdRaw> {
        let path = self.urpd_dir(cohort)?.join(date.to_string());

        if !path.exists() {
            return Err(Error::NotFound(format!(
                "No URPD for cohort '{cohort}' on {date}"
            )));
        }

        UrpdRaw::deserialize(&fs::read(&path)?)
    }

    /// URPD for a cohort on a specific date.
    pub fn urpd_at(&self, cohort: &Cohort, date: Date, agg: UrpdAggregation) -> Result<Urpd> {
        let raw = self.urpd_raw(cohort, date)?;
        let day1 = Day1::try_from(date)?;
        let close = self
            .computer()
            .prices
            .split
            .close
            .cents
            .day1
            .collect_one_flat(day1)
            .ok_or_else(|| Error::NotFound(format!("No price data for {date}")))?;
        Ok(Urpd::build(cohort.clone(), date, close, &raw, agg))
    }

    /// URPD for the most recently available date in a cohort.
    pub fn urpd_latest(&self, cohort: &Cohort, agg: UrpdAggregation) -> Result<Urpd> {
        let dates = self.urpd_dates(cohort)?;
        let date = *dates
            .last()
            .ok_or_else(|| Error::NotFound(format!("No URPD available for cohort '{cohort}'")))?;
        self.urpd_at(cohort, date, agg)
    }
}