zen_expression/functions/
date_method.rs

1use crate::functions::defs::{
2    CompositeFunction, FunctionDefinition, FunctionSignature, StaticFunction,
3};
4use crate::vm::date::DurationUnit;
5use std::rc::Rc;
6use strum_macros::{Display, EnumIter, EnumString, IntoStaticStr};
7
8#[derive(Debug, PartialEq, Eq, Hash, Display, EnumString, EnumIter, IntoStaticStr, Clone, Copy)]
9#[strum(serialize_all = "camelCase")]
10pub enum DateMethod {
11    Add,
12    Sub,
13    Set,
14    Format,
15    StartOf,
16    EndOf,
17    Diff,
18    Tz,
19
20    // Compare
21    IsSame,
22    IsBefore,
23    IsAfter,
24    IsSameOrBefore,
25    IsSameOrAfter,
26
27    // Getters
28    Second,
29    Minute,
30    Hour,
31    Day,
32    DayOfYear,
33    Week,
34    Weekday,
35    Month,
36    Quarter,
37    Year,
38    Timestamp,
39    OffsetName,
40
41    IsValid,
42    IsYesterday,
43    IsToday,
44    IsTomorrow,
45    IsLeapYear,
46}
47
48enum CompareOperation {
49    IsSame,
50    IsBefore,
51    IsAfter,
52    IsSameOrBefore,
53    IsSameOrAfter,
54}
55
56enum GetterOperation {
57    Second,
58    Minute,
59    Hour,
60    Day,
61    Weekday,
62    DayOfYear,
63    Week,
64    Month,
65    Quarter,
66    Year,
67    Timestamp,
68    OffsetName,
69
70    IsValid,
71    IsYesterday,
72    IsToday,
73    IsTomorrow,
74    IsLeapYear,
75}
76
77impl From<&DateMethod> for Rc<dyn FunctionDefinition> {
78    fn from(value: &DateMethod) -> Self {
79        use crate::variable::VariableType as VT;
80        use DateMethod as DM;
81
82        let unit_vt = DurationUnit::variable_type();
83
84        let op_signature = vec![
85            FunctionSignature {
86                parameters: vec![VT::Date, VT::String],
87                return_type: VT::Date,
88            },
89            FunctionSignature {
90                parameters: vec![VT::Date, VT::Number, unit_vt.clone()],
91                return_type: VT::Date,
92            },
93        ];
94
95        match value {
96            DM::Add => Rc::new(CompositeFunction {
97                implementation: Rc::new(imp::add),
98                signatures: op_signature.clone(),
99            }),
100            DM::Sub => Rc::new(CompositeFunction {
101                implementation: Rc::new(imp::sub),
102                signatures: op_signature.clone(),
103            }),
104            DM::Set => Rc::new(StaticFunction {
105                implementation: Rc::new(imp::set),
106                signature: FunctionSignature {
107                    parameters: vec![VT::Date, unit_vt.clone(), VT::Number],
108                    return_type: VT::Date,
109                },
110            }),
111            DM::Tz => Rc::new(StaticFunction {
112                implementation: Rc::new(imp::tz),
113                signature: FunctionSignature {
114                    parameters: vec![VT::Date, VT::String],
115                    return_type: VT::Date,
116                },
117            }),
118            DM::Format => Rc::new(CompositeFunction {
119                implementation: Rc::new(imp::format),
120                signatures: vec![
121                    FunctionSignature {
122                        parameters: vec![VT::Date],
123                        return_type: VT::String,
124                    },
125                    FunctionSignature {
126                        parameters: vec![VT::Date, VT::String],
127                        return_type: VT::String,
128                    },
129                ],
130            }),
131            DM::StartOf => Rc::new(StaticFunction {
132                implementation: Rc::new(imp::start_of),
133                signature: FunctionSignature {
134                    parameters: vec![VT::Date, unit_vt.clone()],
135                    return_type: VT::Date,
136                },
137            }),
138            DM::EndOf => Rc::new(StaticFunction {
139                implementation: Rc::new(imp::end_of),
140                signature: FunctionSignature {
141                    parameters: vec![VT::Date, unit_vt.clone()],
142                    return_type: VT::Date,
143                },
144            }),
145            DM::Diff => Rc::new(CompositeFunction {
146                implementation: Rc::new(imp::diff),
147                signatures: vec![
148                    FunctionSignature {
149                        parameters: vec![VT::Date, VT::Date],
150                        return_type: VT::Number,
151                    },
152                    FunctionSignature {
153                        parameters: vec![VT::Date, VT::Date, unit_vt.clone()],
154                        return_type: VT::Number,
155                    },
156                ],
157            }),
158            DateMethod::IsSame => imp::compare_using(CompareOperation::IsSame),
159            DateMethod::IsBefore => imp::compare_using(CompareOperation::IsBefore),
160            DateMethod::IsAfter => imp::compare_using(CompareOperation::IsAfter),
161            DateMethod::IsSameOrBefore => imp::compare_using(CompareOperation::IsSameOrBefore),
162            DateMethod::IsSameOrAfter => imp::compare_using(CompareOperation::IsSameOrAfter),
163
164            DateMethod::Second => imp::getter(GetterOperation::Second),
165            DateMethod::Minute => imp::getter(GetterOperation::Minute),
166            DateMethod::Hour => imp::getter(GetterOperation::Hour),
167            DateMethod::Day => imp::getter(GetterOperation::Day),
168            DateMethod::Weekday => imp::getter(GetterOperation::Weekday),
169            DateMethod::DayOfYear => imp::getter(GetterOperation::DayOfYear),
170            DateMethod::Week => imp::getter(GetterOperation::Week),
171            DateMethod::Month => imp::getter(GetterOperation::Month),
172            DateMethod::Quarter => imp::getter(GetterOperation::Quarter),
173            DateMethod::Year => imp::getter(GetterOperation::Year),
174            DateMethod::Timestamp => imp::getter(GetterOperation::Timestamp),
175            DateMethod::OffsetName => imp::getter(GetterOperation::OffsetName),
176
177            DateMethod::IsValid => imp::getter(GetterOperation::IsValid),
178            DateMethod::IsYesterday => imp::getter(GetterOperation::IsYesterday),
179            DateMethod::IsToday => imp::getter(GetterOperation::IsToday),
180            DateMethod::IsTomorrow => imp::getter(GetterOperation::IsTomorrow),
181            DateMethod::IsLeapYear => imp::getter(GetterOperation::IsLeapYear),
182        }
183    }
184}
185
186mod imp {
187    use crate::functions::arguments::Arguments;
188    use crate::functions::date_method::{CompareOperation, GetterOperation};
189    use crate::functions::defs::{
190        CompositeFunction, FunctionDefinition, FunctionSignature, StaticFunction,
191    };
192    use crate::variable::VariableType as VT;
193    use crate::vm::date::{Duration, DurationUnit};
194    use crate::vm::VmDate;
195    use crate::Variable as V;
196    use anyhow::{anyhow, Context};
197    use chrono::{Datelike, Timelike};
198    use chrono_tz::Tz;
199    use rust_decimal::prelude::{FromPrimitive, ToPrimitive};
200    use rust_decimal::Decimal;
201    use std::rc::Rc;
202    use std::str::FromStr;
203
204    fn __internal_extract_duration(args: &Arguments, from: usize) -> anyhow::Result<Duration> {
205        match args.var(from)? {
206            V::String(s) => Ok(Duration::parse(s.as_ref())?),
207            V::Number(n) => {
208                let unit = __internal_extract_duration_unit(args, from + 1)?;
209                Ok(Duration::from_unit(*n, unit).context("Invalid duration unit")?)
210            }
211            _ => Err(anyhow!("Invalid duration arguments")),
212        }
213    }
214
215    fn __internal_extract_duration_unit(
216        args: &Arguments,
217        pos: usize,
218    ) -> anyhow::Result<DurationUnit> {
219        let unit_str = args.str(pos)?;
220        DurationUnit::parse(unit_str).context("Invalid duration unit")
221    }
222
223    fn __internal_extract_duration_unit_opt(
224        args: &Arguments,
225        pos: usize,
226    ) -> anyhow::Result<Option<DurationUnit>> {
227        let unit_ostr = args.ostr(pos)?;
228        let Some(unit_str) = unit_ostr else {
229            return Ok(None);
230        };
231
232        Ok(Some(
233            DurationUnit::parse(unit_str).context("Invalid duration unit")?,
234        ))
235    }
236
237    pub fn add(args: Arguments) -> anyhow::Result<V> {
238        let this = args.dynamic::<VmDate>(0)?;
239        let duration = __internal_extract_duration(&args, 1)?;
240
241        let date_time = this.add(duration);
242        Ok(V::Dynamic(Rc::new(date_time)))
243    }
244
245    pub fn sub(args: Arguments) -> anyhow::Result<V> {
246        let this = args.dynamic::<VmDate>(0)?;
247        let duration = __internal_extract_duration(&args, 1)?;
248
249        let date_time = this.sub(duration);
250        Ok(V::Dynamic(Rc::new(date_time)))
251    }
252
253    pub fn set(args: Arguments) -> anyhow::Result<V> {
254        let this = args.dynamic::<VmDate>(0)?;
255        let unit = __internal_extract_duration_unit(&args, 1)?;
256        let value = args.number(2)?;
257
258        let value_u32 = value.to_u32().context("Invalid duration value")?;
259
260        let date_time = this.set(value_u32, unit);
261        Ok(V::Dynamic(Rc::new(date_time)))
262    }
263
264    pub fn format(args: Arguments) -> anyhow::Result<V> {
265        let this = args.dynamic::<VmDate>(0)?;
266        let format = args.ostr(1)?;
267
268        let formatted = this.format(format);
269        Ok(V::String(Rc::from(formatted)))
270    }
271
272    pub fn start_of(args: Arguments) -> anyhow::Result<V> {
273        let this = args.dynamic::<VmDate>(0)?;
274        let unit = __internal_extract_duration_unit(&args, 1)?;
275
276        let date_time = this.start_of(unit);
277        Ok(V::Dynamic(Rc::new(date_time)))
278    }
279
280    pub fn end_of(args: Arguments) -> anyhow::Result<V> {
281        let this = args.dynamic::<VmDate>(0)?;
282        let unit = __internal_extract_duration_unit(&args, 1)?;
283
284        let date_time = this.end_of(unit);
285        Ok(V::Dynamic(Rc::new(date_time)))
286    }
287
288    pub fn diff(args: Arguments) -> anyhow::Result<V> {
289        let this = args.dynamic::<VmDate>(0)?;
290        let date_time = VmDate::new(args.var(1)?.clone(), None);
291        let maybe_unit = __internal_extract_duration_unit_opt(&args, 2)?;
292
293        let var = match this
294            .diff(&date_time, maybe_unit)
295            .and_then(Decimal::from_i64)
296        {
297            Some(n) => V::Number(n),
298            None => V::Null,
299        };
300
301        Ok(var)
302    }
303
304    pub fn tz(args: Arguments) -> anyhow::Result<V> {
305        let this = args.dynamic::<VmDate>(0)?;
306        let tz_str = args.str(1)?;
307
308        let timezone = Tz::from_str(tz_str).context("Invalid timezone")?;
309        Ok(V::Dynamic(Rc::new(this.tz(timezone))))
310    }
311
312    pub fn compare_using(op: CompareOperation) -> Rc<dyn FunctionDefinition> {
313        Rc::new(CompositeFunction {
314            signatures: vec![
315                FunctionSignature {
316                    parameters: vec![VT::Date, VT::Date],
317                    return_type: VT::Date,
318                },
319                FunctionSignature {
320                    parameters: vec![VT::Date, VT::Date, DurationUnit::variable_type()],
321                    return_type: VT::Date,
322                },
323            ],
324            implementation: Rc::new(move |args: Arguments| -> anyhow::Result<V> {
325                let this = args.dynamic::<VmDate>(0)?;
326                let date_time = VmDate::new(args.var(1)?.clone(), None);
327                let maybe_unit = __internal_extract_duration_unit_opt(&args, 2)?;
328
329                let check = match op {
330                    CompareOperation::IsSame => this.is_same(&date_time, maybe_unit),
331                    CompareOperation::IsBefore => this.is_before(&date_time, maybe_unit),
332                    CompareOperation::IsAfter => this.is_after(&date_time, maybe_unit),
333                    CompareOperation::IsSameOrBefore => {
334                        this.is_same_or_before(&date_time, maybe_unit)
335                    }
336                    CompareOperation::IsSameOrAfter => {
337                        this.is_same_or_after(&date_time, maybe_unit)
338                    }
339                };
340
341                Ok(V::Bool(check))
342            }),
343        })
344    }
345
346    pub fn getter(op: GetterOperation) -> Rc<dyn FunctionDefinition> {
347        Rc::new(StaticFunction {
348            signature: FunctionSignature {
349                parameters: vec![VT::Date],
350                return_type: match op {
351                    GetterOperation::Second
352                    | GetterOperation::Minute
353                    | GetterOperation::Hour
354                    | GetterOperation::Day
355                    | GetterOperation::Weekday
356                    | GetterOperation::DayOfYear
357                    | GetterOperation::Week
358                    | GetterOperation::Month
359                    | GetterOperation::Quarter
360                    | GetterOperation::Year
361                    | GetterOperation::Timestamp => VT::Number,
362                    GetterOperation::IsValid
363                    | GetterOperation::IsYesterday
364                    | GetterOperation::IsToday
365                    | GetterOperation::IsTomorrow
366                    | GetterOperation::IsLeapYear => VT::Bool,
367                    GetterOperation::OffsetName => VT::String,
368                },
369            },
370            implementation: Rc::new(move |args: Arguments| -> anyhow::Result<V> {
371                let this = args.dynamic::<VmDate>(0)?;
372                if let GetterOperation::IsValid = op {
373                    return Ok(V::Bool(this.is_valid()));
374                }
375
376                let Some(dt) = this.0 else {
377                    return Ok(V::Null);
378                };
379
380                Ok(match op {
381                    GetterOperation::Second => V::Number(dt.second().into()),
382                    GetterOperation::Minute => V::Number(dt.minute().into()),
383                    GetterOperation::Hour => V::Number(dt.hour().into()),
384                    GetterOperation::Day => V::Number(dt.day().into()),
385                    GetterOperation::Weekday => V::Number(dt.weekday().number_from_monday().into()),
386                    GetterOperation::DayOfYear => V::Number(dt.ordinal().into()),
387                    GetterOperation::Week => V::Number(dt.iso_week().week().into()),
388                    GetterOperation::Month => V::Number(dt.month().into()),
389                    GetterOperation::Quarter => V::Number(dt.quarter().into()),
390                    GetterOperation::Year => V::Number(dt.year().into()),
391                    GetterOperation::Timestamp => V::Number(dt.timestamp_millis().into()),
392                    // Boolean
393                    GetterOperation::IsValid => V::Bool(true),
394                    GetterOperation::IsYesterday => {
395                        V::Bool(this.is_same(&VmDate::yesterday(), Some(DurationUnit::Day)))
396                    }
397                    GetterOperation::IsToday => {
398                        V::Bool(this.is_same(&VmDate::now(), Some(DurationUnit::Day)))
399                    }
400                    GetterOperation::IsTomorrow => {
401                        V::Bool(this.is_same(&VmDate::tomorrow(), Some(DurationUnit::Day)))
402                    }
403                    GetterOperation::IsLeapYear => V::Bool(dt.date_naive().leap_year()),
404                    // String
405                    GetterOperation::OffsetName => V::String(Rc::from(dt.timezone().name())),
406                })
407            }),
408        })
409    }
410}