progscrape_application/persist/
shard.rs

1use std::{fmt::Debug, ops::RangeInclusive};
2
3use progscrape_scrapers::StoryDate;
4use serde::{Deserialize, Serialize};
5
6#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
7pub struct Shard(u16);
8
9impl ToString for Shard {
10    fn to_string(&self) -> String {
11        format!("{:04}-{:02}", self.0 / 12, self.0 % 12 + 1)
12    }
13}
14
15impl Debug for Shard {
16    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
17        f.write_fmt(format_args!("{:04}-{:02}", self.0 / 12, self.0 % 12 + 1))
18    }
19}
20
21impl Default for Shard {
22    fn default() -> Self {
23        Shard::from_year_month(2000, 1)
24    }
25}
26
27impl std::ops::Add<u16> for Shard {
28    type Output = Shard;
29    fn add(self, rhs: u16) -> Self::Output {
30        Shard(self.0 + rhs)
31    }
32}
33
34impl std::ops::Sub<u16> for Shard {
35    type Output = Shard;
36    fn sub(self, rhs: u16) -> Self::Output {
37        Shard(self.0 - rhs)
38    }
39}
40
41#[derive(PartialEq, Eq, PartialOrd, Ord)]
42pub enum ShardOrder {
43    OldestFirst,
44    NewestFirst,
45}
46
47#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
48pub struct ShardRange {
49    range: Option<(Shard, Shard)>,
50}
51
52impl ShardRange {
53    pub fn new() -> Self {
54        Default::default()
55    }
56
57    pub fn new_from(range: RangeInclusive<Shard>) -> Self {
58        Self {
59            range: Some((*range.start(), *range.end())),
60        }
61    }
62
63    pub fn iterate(&self, order: ShardOrder) -> impl Iterator<Item = Shard> {
64        let (mut start, end) = self.range.unwrap_or((Shard::MAX, Shard::MIN));
65        let orig_start = start;
66        std::iter::from_fn(move || {
67            if start > end {
68                None
69            } else {
70                let next = Some(if order == ShardOrder::OldestFirst {
71                    start
72                } else {
73                    Shard((end.0 - start.0) + orig_start.0)
74                });
75                start = start + 1;
76                next
77            }
78        })
79    }
80
81    pub fn include(&mut self, shard: Shard) {
82        if let Some(range) = self.range {
83            self.range = Some((range.0.min(shard), range.1.max(shard)))
84        } else {
85            self.range = Some((shard, shard))
86        }
87    }
88}
89
90impl Shard {
91    pub const MIN: Shard = Shard(u16::MIN);
92    pub const MAX: Shard = Shard(u16::MAX);
93
94    pub fn from_year_month(year: u16, month: u8) -> Self {
95        assert!(month > 0);
96        Shard(year * 12 + month as u16 - 1)
97    }
98
99    pub fn from_string(s: &str) -> Option<Self> {
100        if let Some((a, b)) = s.split_once('-') {
101            if let (Ok(a), Ok(b)) = (str::parse(a), str::parse(b)) {
102                return Some(Self::from_year_month(a, b));
103            }
104        }
105        None
106    }
107
108    pub fn from_date_time(date: StoryDate) -> Self {
109        Self::from_year_month(date.year() as u16, date.month() as u8)
110    }
111
112    pub fn plus_months(&self, months: i8) -> Self {
113        let ordinal = self.0 as i16 + months as i16;
114        Self(ordinal as u16)
115    }
116
117    pub fn sub_months(&self, months: i8) -> Self {
118        self.plus_months(-months)
119    }
120}
121
122#[cfg(test)]
123mod test {
124    use itertools::Itertools;
125
126    use super::*;
127
128    #[test]
129    fn test_shard_iterator() {
130        let range = ShardRange::new_from(
131            Shard::from_year_month(2000, 1)..=Shard::from_year_month(2000, 12),
132        );
133        assert_eq!(range.iterate(ShardOrder::OldestFirst).count(), 12);
134        assert_eq!(range.iterate(ShardOrder::NewestFirst).count(), 12);
135
136        let in_order = range.iterate(ShardOrder::OldestFirst).collect_vec();
137        let mut rev_order = range.iterate(ShardOrder::NewestFirst).collect_vec();
138        rev_order.reverse();
139
140        assert_eq!(in_order, rev_order);
141    }
142}