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