hegel/extras/chrono/
generators.rs

1use crate::cbor_utils::{cbor_array, cbor_map};
2use crate::generators::{BasicGenerator, DefaultGenerator, Generator, TestCase};
3use chrono::{
4    DateTime, Datelike, FixedOffset, NaiveDate, NaiveDateTime, NaiveTime, NaiveWeek, TimeDelta,
5    TimeZone, Timelike, Utc, Weekday, WeekdaySet,
6};
7use ciborium::Value;
8use std::marker::PhantomData;
9
10/// Convert a [`NaiveTime`] to a total nanosecond count from midnight.
11///
12/// chrono encodes leap seconds by letting `nanosecond()` reach into
13/// `[1_000_000_000, 2_000_000_000)`; those bits are preserved here so the
14/// conversion is lossless and round-trips through [`total_nanos_to_time`].
15fn time_to_total_nanos(t: NaiveTime) -> i64 {
16    let secs = i64::from(t.num_seconds_from_midnight());
17    let nanos = i64::from(t.nanosecond());
18    secs * 1_000_000_000 + nanos
19}
20
21/// Inverse of [`time_to_total_nanos`].
22fn total_nanos_to_time(total: i64) -> NaiveTime {
23    // For leap-second values, chrono keeps secs at 86399 and pushes the extra
24    // second into the nanos field. Detect that range and reconstruct in kind.
25    let (secs, nanos) = if total >= 86_400 * 1_000_000_000 {
26        (86_399, (total - 86_399 * 1_000_000_000) as u32)
27    } else {
28        (
29            (total / 1_000_000_000) as u32,
30            (total % 1_000_000_000) as u32,
31        )
32    };
33    NaiveTime::from_num_seconds_from_midnight_opt(secs, nanos).unwrap()
34}
35
36/// Default upper bound for [`naive_times`]. Excludes leap-second representation
37/// — users who want to generate leap seconds must opt in via [`NaiveTimeGenerator::max_value`].
38fn naive_time_default_max() -> NaiveTime {
39    NaiveTime::from_hms_nano_opt(23, 59, 59, 999_999_999).unwrap()
40}
41
42/// Generator for [`chrono::WeekdaySet`] values. Created by [`weekday_sets()`].
43pub struct WeekdaySetGenerator;
44
45impl WeekdaySetGenerator {
46    fn build_schema(&self) -> Value {
47        cbor_map! {
48            "type" => "list",
49            "unique" => true,
50            "elements" => cbor_map! {
51                "type" => "integer",
52                "min_value" => 0u64,
53                "max_value" => 6u64,
54            },
55            "min_size" => 0u64,
56            "max_size" => 7u64,
57        }
58    }
59}
60
61impl Generator<WeekdaySet> for WeekdaySetGenerator {
62    fn as_basic(&self) -> Option<BasicGenerator<'_, WeekdaySet>> {
63        Some(BasicGenerator::new(self.build_schema(), |raw| {
64            let arr = raw.into_array().unwrap();
65            let mut set = WeekdaySet::EMPTY;
66            for v in arr {
67                let n: u8 = crate::generators::deserialize_value(v);
68                set.insert(Weekday::try_from(n).unwrap());
69            }
70            set
71        }))
72    }
73}
74
75/// Generate [`chrono::WeekdaySet`] values.
76///
77/// # Example
78///
79/// ```no_run
80/// use hegel::extras::chrono as chrono_gs;
81///
82/// #[hegel::test]
83/// fn my_test(tc: hegel::TestCase) {
84///     let s: chrono::WeekdaySet = tc.draw(chrono_gs::weekday_sets());
85///     assert!(s.len() <= 7);
86/// }
87/// ```
88pub fn weekday_sets() -> WeekdaySetGenerator {
89    WeekdaySetGenerator
90}
91
92/// Generator for [`chrono::FixedOffset`] values. Created by [`fixed_offsets()`].
93pub struct FixedOffsetGenerator {
94    min_value: FixedOffset,
95    max_value: FixedOffset,
96}
97
98impl FixedOffsetGenerator {
99    /// Set the minimum offset (inclusive).
100    pub fn min_value(mut self, min: FixedOffset) -> Self {
101        self.min_value = min;
102        self
103    }
104
105    /// Set the maximum offset (inclusive).
106    pub fn max_value(mut self, max: FixedOffset) -> Self {
107        self.max_value = max;
108        self
109    }
110
111    fn build_schema(&self) -> Value {
112        let min_secs = self.min_value.local_minus_utc();
113        let max_secs = self.max_value.local_minus_utc();
114        assert!(min_secs <= max_secs, "Cannot have max_value < min_value");
115        cbor_map! {
116            "type" => "integer",
117            "min_value" => i64::from(min_secs),
118            "max_value" => i64::from(max_secs),
119        }
120    }
121}
122
123impl Generator<FixedOffset> for FixedOffsetGenerator {
124    fn as_basic(&self) -> Option<BasicGenerator<'_, FixedOffset>> {
125        Some(BasicGenerator::new(self.build_schema(), |raw| {
126            let secs: i32 = crate::generators::deserialize_value(raw);
127            FixedOffset::east_opt(secs).unwrap()
128        }))
129    }
130}
131
132/// Generate [`chrono::FixedOffset`] values.
133///
134/// # Example
135///
136/// ```no_run
137/// use chrono::FixedOffset;
138/// use hegel::extras::chrono as chrono_gs;
139///
140/// #[hegel::test]
141/// fn my_test(tc: hegel::TestCase) {
142///     let max = FixedOffset::east_opt(12 * 3600).unwrap();
143///     let off = tc.draw(chrono_gs::fixed_offsets().max_value(max));
144///     assert!(off.local_minus_utc() <= 12 * 3600);
145/// }
146/// ```
147pub fn fixed_offsets() -> FixedOffsetGenerator {
148    FixedOffsetGenerator {
149        min_value: FixedOffset::west_opt(86_399).unwrap(),
150        max_value: FixedOffset::east_opt(86_399).unwrap(),
151    }
152}
153
154/// Convert a [`TimeDelta`] to a total nanosecond count.
155///
156/// `TimeDelta` is internally `(secs: i64, nanos: u32 in [0, 1e9))`. The total
157/// nanosecond magnitude exceeds `i64`'s range by a factor of ~10^9, so we
158/// widen to `i128` (which has ~10^10 of headroom past `TimeDelta::MAX`).
159fn timedelta_to_nanos(d: TimeDelta) -> i128 {
160    i128::from(d.num_seconds()) * 1_000_000_000 + i128::from(d.subsec_nanos())
161}
162
163/// Inverse of [`timedelta_to_nanos`].
164fn nanos_to_timedelta(n: i128) -> TimeDelta {
165    let secs = n.div_euclid(1_000_000_000) as i64;
166    let nanos = n.rem_euclid(1_000_000_000) as u32;
167    TimeDelta::new(secs, nanos).unwrap()
168}
169
170/// Generator for [`chrono::TimeDelta`] values. Created by [`time_deltas()`].
171pub struct TimeDeltaGenerator {
172    min_value: TimeDelta,
173    max_value: TimeDelta,
174}
175
176impl TimeDeltaGenerator {
177    /// Set the minimum delta (inclusive).
178    pub fn min_value(mut self, min: TimeDelta) -> Self {
179        self.min_value = min;
180        self
181    }
182
183    /// Set the maximum delta (inclusive).
184    pub fn max_value(mut self, max: TimeDelta) -> Self {
185        self.max_value = max;
186        self
187    }
188
189    fn build_schema(&self) -> Value {
190        assert!(
191            self.min_value <= self.max_value,
192            "Cannot have max_value < min_value"
193        );
194        cbor_map! {
195            "type" => "integer",
196            "min_value" => timedelta_to_nanos(self.min_value),
197            "max_value" => timedelta_to_nanos(self.max_value),
198        }
199    }
200}
201
202impl Generator<TimeDelta> for TimeDeltaGenerator {
203    fn as_basic(&self) -> Option<BasicGenerator<'_, TimeDelta>> {
204        Some(BasicGenerator::new(self.build_schema(), |raw| {
205            let nanos: i128 = crate::generators::deserialize_value(raw);
206            nanos_to_timedelta(nanos)
207        }))
208    }
209}
210
211/// Generate [`chrono::TimeDelta`] values.
212///
213/// Defaults span the full `TimeDelta::MIN..=TimeDelta::MAX` range. Use the
214/// builder methods to constrain.
215///
216/// # Example
217///
218/// ```no_run
219/// use chrono::TimeDelta;
220/// use hegel::extras::chrono as chrono_gs;
221///
222/// #[hegel::test]
223/// fn my_test(tc: hegel::TestCase) {
224///     let max = TimeDelta::seconds(60);
225///     let d = tc.draw(chrono_gs::time_deltas().min_value(TimeDelta::zero()).max_value(max));
226///     assert!(d >= TimeDelta::zero() && d <= max);
227/// }
228/// ```
229pub fn time_deltas() -> TimeDeltaGenerator {
230    TimeDeltaGenerator {
231        min_value: TimeDelta::MIN,
232        max_value: TimeDelta::MAX,
233    }
234}
235
236/// Generator for [`chrono::NaiveDate`] values. Created by [`naive_dates()`].
237///
238/// Internally the date is generated as a count of days from the Common Era
239/// epoch (`NaiveDate::from_num_days_from_ce_opt`).
240pub struct NaiveDateGenerator {
241    min_value: NaiveDate,
242    max_value: NaiveDate,
243}
244
245impl NaiveDateGenerator {
246    /// Set the minimum date (inclusive).
247    pub fn min_value(mut self, min: NaiveDate) -> Self {
248        self.min_value = min;
249        self
250    }
251
252    /// Set the maximum date (inclusive).
253    pub fn max_value(mut self, max: NaiveDate) -> Self {
254        self.max_value = max;
255        self
256    }
257
258    fn build_schema(&self) -> Value {
259        assert!(
260            self.min_value <= self.max_value,
261            "Cannot have max_value < min_value"
262        );
263        cbor_map! {
264            "type" => "integer",
265            "min_value" => self.min_value.num_days_from_ce(),
266            "max_value" => self.max_value.num_days_from_ce(),
267        }
268    }
269}
270
271impl Generator<NaiveDate> for NaiveDateGenerator {
272    fn as_basic(&self) -> Option<BasicGenerator<'_, NaiveDate>> {
273        Some(BasicGenerator::new(self.build_schema(), |raw| {
274            let n: i32 = crate::generators::deserialize_value(raw);
275            NaiveDate::from_num_days_from_ce_opt(n).unwrap()
276        }))
277    }
278}
279
280/// Generate [`chrono::NaiveDate`] values.
281///
282/// # Example
283///
284/// ```no_run
285/// use chrono::NaiveDate;
286/// use hegel::extras::chrono as chrono_gs;
287///
288/// #[hegel::test]
289/// fn my_test(tc: hegel::TestCase) {
290///     let min = NaiveDate::from_ymd_opt(2024, 1, 1).unwrap();
291///     let max = NaiveDate::from_ymd_opt(2024, 12, 31).unwrap();
292///     let d = tc.draw(chrono_gs::naive_dates().min_value(min).max_value(max));
293///     assert_eq!(d.iso_week().year(), 2024);
294/// }
295/// ```
296pub fn naive_dates() -> NaiveDateGenerator {
297    NaiveDateGenerator {
298        min_value: NaiveDate::MIN,
299        max_value: NaiveDate::MAX,
300    }
301}
302
303/// Generator for [`chrono::NaiveTime`] values. Created by [`naive_times()`].
304pub struct NaiveTimeGenerator {
305    min_value: NaiveTime,
306    max_value: NaiveTime,
307}
308
309impl NaiveTimeGenerator {
310    /// Set the minimum time (inclusive).
311    pub fn min_value(mut self, min: NaiveTime) -> Self {
312        self.min_value = min;
313        self
314    }
315
316    /// Set the maximum time (inclusive).
317    pub fn max_value(mut self, max: NaiveTime) -> Self {
318        self.max_value = max;
319        self
320    }
321
322    fn build_schema(&self) -> Value {
323        let min_nanos = time_to_total_nanos(self.min_value);
324        let max_nanos = time_to_total_nanos(self.max_value);
325        assert!(min_nanos <= max_nanos, "Cannot have max_value < min_value");
326        cbor_map! {
327            "type" => "integer",
328            "min_value" => min_nanos,
329            "max_value" => max_nanos,
330        }
331    }
332}
333
334impl Generator<NaiveTime> for NaiveTimeGenerator {
335    fn as_basic(&self) -> Option<BasicGenerator<'_, NaiveTime>> {
336        Some(BasicGenerator::new(self.build_schema(), |raw| {
337            let n: i64 = crate::generators::deserialize_value(raw);
338            total_nanos_to_time(n)
339        }))
340    }
341}
342
343/// Generate [`chrono::NaiveTime`] values.
344///
345/// # Example
346///
347/// ```no_run
348/// use chrono::NaiveTime;
349/// use hegel::extras::chrono as chrono_gs;
350///
351/// #[hegel::test]
352/// fn my_test(tc: hegel::TestCase) {
353///     let max = NaiveTime::from_hms_opt(12, 0, 0).unwrap();
354///     let t = tc.draw(chrono_gs::naive_times().max_value(max));
355///     assert!(t <= max);
356/// }
357/// ```
358pub fn naive_times() -> NaiveTimeGenerator {
359    NaiveTimeGenerator {
360        min_value: NaiveTime::MIN,
361        max_value: naive_time_default_max(),
362    }
363}
364
365/// Generator for [`chrono::NaiveDateTime`] values. Created by [`naive_datetimes()`].
366pub struct NaiveDateTimeGenerator {
367    min_value: NaiveDateTime,
368    max_value: NaiveDateTime,
369}
370
371impl NaiveDateTimeGenerator {
372    /// Set the minimum datetime (inclusive).
373    pub fn min_value(mut self, min: NaiveDateTime) -> Self {
374        self.min_value = min;
375        self
376    }
377
378    /// Set the maximum datetime (inclusive).
379    pub fn max_value(mut self, max: NaiveDateTime) -> Self {
380        self.max_value = max;
381        self
382    }
383}
384
385impl Generator<NaiveDateTime> for NaiveDateTimeGenerator {
386    fn as_basic(&self) -> Option<BasicGenerator<'_, NaiveDateTime>> {
387        assert!(
388            self.min_value <= self.max_value,
389            "Cannot have max_value < min_value"
390        );
391        let schema = cbor_map! {
392            "type" => "integer",
393            "min_value" => datetime_to_nanos(&self.min_value.and_utc()),
394            "max_value" => datetime_to_nanos(&self.max_value.and_utc()),
395        };
396        Some(BasicGenerator::new(schema, |raw| {
397            let n: i128 = crate::generators::deserialize_value(raw);
398            nanos_to_utc_datetime(n).naive_utc()
399        }))
400    }
401}
402
403/// Generate [`chrono::NaiveDateTime`] values.
404///
405/// Defaults span chrono's full `DateTime<Utc>` window —
406/// [`DateTime::<Utc>::MIN_UTC`] through [`DateTime::<Utc>::MAX_UTC`].
407///
408/// # Example
409///
410/// ```no_run
411/// use chrono::{NaiveDate, NaiveDateTime, NaiveTime};
412/// use hegel::extras::chrono as chrono_gs;
413///
414/// #[hegel::test]
415/// fn my_test(tc: hegel::TestCase) {
416///     let min = NaiveDateTime::new(
417///         NaiveDate::from_ymd_opt(2024, 1, 1).unwrap(),
418///         NaiveTime::from_hms_opt(0, 0, 0).unwrap(),
419///     );
420///     let max = NaiveDateTime::new(
421///         NaiveDate::from_ymd_opt(2024, 12, 31).unwrap(),
422///         NaiveTime::from_hms_nano_opt(23, 59, 59, 999_999_999).unwrap(),
423///     );
424///     let dt = tc.draw(chrono_gs::naive_datetimes().min_value(min).max_value(max));
425///     assert!(dt >= min && dt <= max);
426/// }
427/// ```
428pub fn naive_datetimes() -> NaiveDateTimeGenerator {
429    NaiveDateTimeGenerator {
430        min_value: DateTime::<Utc>::MIN_UTC.naive_utc(),
431        max_value: DateTime::<Utc>::MAX_UTC.naive_utc(),
432    }
433}
434
435/// Generator for [`chrono::NaiveWeek`] values. Created by [`naive_weeks()`].
436pub struct NaiveWeekGenerator<S = <Weekday as DefaultGenerator>::Generator> {
437    date_gen: NaiveDateGenerator,
438    start_gen: S,
439}
440
441impl<S> NaiveWeekGenerator<S> {
442    /// Set the minimum source date (inclusive).
443    pub fn min_date(mut self, min: NaiveDate) -> Self {
444        self.date_gen = self.date_gen.min_value(min);
445        self
446    }
447
448    /// Set the maximum source date (inclusive).
449    pub fn max_date(mut self, max: NaiveDate) -> Self {
450        self.date_gen = self.date_gen.max_value(max);
451        self
452    }
453
454    /// Replace the start-of-week generator.
455    pub fn weekday_starts<S2>(self, start_gen: S2) -> NaiveWeekGenerator<S2>
456    where
457        S2: Generator<Weekday>,
458    {
459        NaiveWeekGenerator {
460            date_gen: self.date_gen,
461            start_gen,
462        }
463    }
464}
465
466impl<S: Generator<Weekday>> Generator<NaiveWeek> for NaiveWeekGenerator<S> {
467    fn as_basic(&self) -> Option<BasicGenerator<'_, NaiveWeek>> {
468        let date_basic = self.date_gen.as_basic()?;
469        let start_basic = self.start_gen.as_basic()?;
470        let schema = cbor_map! {
471            "type" => "tuple",
472            "elements" => cbor_array![
473                date_basic.schema().clone(),
474                start_basic.schema().clone(),
475            ],
476        };
477        Some(BasicGenerator::new(schema, move |raw| {
478            let [d_raw, s_raw]: [Value; 2] = raw.into_array().unwrap().try_into().unwrap();
479            let date = date_basic.parse_raw(d_raw);
480            let start = start_basic.parse_raw(s_raw);
481            date.week(start)
482        }))
483    }
484}
485
486/// Generate [`chrono::NaiveWeek`] values.
487///
488/// Each draw produces a 7-day window around a randomly chosen [`NaiveDate`],
489/// with the start-of-week [`Weekday`] also drawn at random (full 7-day range
490/// by default). Use [`weekday_starts`](NaiveWeekGenerator::weekday_starts) to substitute a
491/// different start-day strategy — e.g. `gs::just(Weekday::Sun)` for Sunday-only.
492///
493/// # Example
494///
495/// ```no_run
496/// use chrono::Weekday;
497/// use hegel::extras::chrono as chrono_gs;
498/// use hegel::generators as gs;
499///
500/// #[hegel::test]
501/// fn my_test(tc: hegel::TestCase) {
502///     // Default: random start day per draw
503///     let _ = tc.draw(chrono_gs::naive_weeks());
504///
505///     // All Sunday-start weeks
506///     let _ = tc.draw(chrono_gs::naive_weeks().weekday_starts(gs::just(Weekday::Sun)));
507/// }
508/// ```
509pub fn naive_weeks() -> NaiveWeekGenerator {
510    NaiveWeekGenerator {
511        date_gen: naive_dates(),
512        start_gen: Weekday::default_generator(),
513    }
514}
515
516/// Convert a [`DateTime`] to a total nanosecond count from the Unix epoch.
517fn datetime_to_nanos<Tz: TimeZone>(dt: &DateTime<Tz>) -> i128 {
518    i128::from(dt.timestamp()) * 1_000_000_000 + i128::from(dt.timestamp_subsec_nanos())
519}
520
521/// Inverse of [`datetime_to_nanos`], producing a `DateTime<Utc>`.
522fn nanos_to_utc_datetime(n: i128) -> DateTime<Utc> {
523    let secs = n.div_euclid(1_000_000_000) as i64;
524    let nsecs = n.rem_euclid(1_000_000_000) as u32;
525    DateTime::<Utc>::from_timestamp(secs, nsecs).unwrap()
526}
527
528/// Generator for [`chrono::DateTime`] values. Created by [`datetimes()`].
529pub struct DateTimeGenerator<G = FixedOffsetGenerator, Tz: TimeZone = FixedOffset> {
530    tz_gen: G,
531    min_value: NaiveDateTime,
532    max_value: NaiveDateTime,
533    _phantom: PhantomData<fn() -> Tz>,
534}
535
536impl<G, Tz: TimeZone> DateTimeGenerator<G, Tz> {
537    /// Set the minimum wall-clock datetime (inclusive).
538    pub fn min_value(mut self, min: NaiveDateTime) -> Self {
539        self.min_value = min;
540        self
541    }
542
543    /// Set the maximum wall-clock datetime (inclusive).
544    pub fn max_value(mut self, max: NaiveDateTime) -> Self {
545        self.max_value = max;
546        self
547    }
548
549    /// Use the specified timezone generator.
550    ///
551    /// The returned generator has output type `DateTime<Tz2>` where `Tz2`
552    /// is the timezone type produced by `tz_gen`. Wall-clock bounds are
553    /// timezone-agnostic and carry across.
554    pub fn timezones<G2, Tz2>(self, tz_gen: G2) -> DateTimeGenerator<G2, Tz2>
555    where
556        G2: Generator<Tz2>,
557        Tz2: TimeZone,
558    {
559        DateTimeGenerator {
560            tz_gen,
561            min_value: self.min_value,
562            max_value: self.max_value,
563            _phantom: PhantomData,
564        }
565    }
566}
567
568impl<G, Tz> Generator<DateTime<Tz>> for DateTimeGenerator<G, Tz>
569where
570    G: Generator<Tz>,
571    Tz: TimeZone + Send + Sync + 'static,
572{
573    fn do_draw(&self, tc: &TestCase) -> DateTime<Tz> {
574        let naive = naive_datetimes()
575            .min_value(self.min_value)
576            .max_value(self.max_value)
577            .do_draw(tc);
578        let tz = self.tz_gen.do_draw(tc);
579        match tz.from_local_datetime(&naive).earliest() {
580            Some(dt) => dt,
581            None => {
582                tc.assume(false);
583                unreachable!()
584            }
585        }
586    }
587}
588
589/// Generate [`chrono::DateTime`] values.
590///
591/// ```no_run
592/// use chrono::{FixedOffset, Utc};
593/// use hegel::extras::chrono as chrono_gs;
594/// use hegel::generators as gs;
595///
596/// // Default: DateTime<FixedOffset> with varying offsets
597/// let g = chrono_gs::datetimes();
598///
599/// // All UTC
600/// let g = chrono_gs::datetimes().timezones(gs::just(Utc));
601///
602/// // All in a single fixed offset
603/// let off = FixedOffset::east_opt(9 * 3600).unwrap();
604/// let g = chrono_gs::datetimes().timezones(gs::just(off));
605///
606/// // Constrain offsets to a sub-range
607/// let g = chrono_gs::datetimes()
608///     .timezones(chrono_gs::fixed_offsets().min_value(off));
609/// ```
610pub fn datetimes() -> DateTimeGenerator<FixedOffsetGenerator, FixedOffset> {
611    DateTimeGenerator {
612        tz_gen: fixed_offsets(),
613        min_value: DateTime::<Utc>::MIN_UTC.naive_utc(),
614        max_value: DateTime::<Utc>::MAX_UTC.naive_utc(),
615        _phantom: PhantomData,
616    }
617}
hegel/extras/chrono/generators.rs

hegel/extras/chrono/
generators.rs
