1use std::str::FromStr;
2
3use ciborium::Value;
4use jiff::civil::{Date, DateTime, Time};
5use jiff::tz::{Offset, TimeZone};
6use jiff::{SignedDuration, Span, Timestamp, Zoned};
7
8use crate::cbor_utils::{cbor_array, cbor_map};
9use crate::generators::{BasicGenerator, BoxedGenerator, Generator, deserialize_value};
10
11pub struct DateGenerator;
13
14impl Generator<Date> for DateGenerator {
15 fn as_basic(&self) -> Option<BasicGenerator<'_, Date>> {
16 Some(BasicGenerator::new(cbor_map! {"type" => "date"}, |raw| {
17 let s: String = deserialize_value(raw);
18 Date::from_str(&s).unwrap()
19 }))
20 }
21}
22
23pub fn dates() -> DateGenerator {
39 DateGenerator
40}
41
42pub struct TimeGenerator;
44
45impl Generator<Time> for TimeGenerator {
46 fn as_basic(&self) -> Option<BasicGenerator<'_, Time>> {
47 Some(BasicGenerator::new(cbor_map! {"type" => "time"}, |raw| {
48 let s: String = deserialize_value(raw);
49 Time::from_str(&s).unwrap()
50 }))
51 }
52}
53
54pub fn times() -> TimeGenerator {
70 TimeGenerator
71}
72
73fn datetime_to_nanos(dt: DateTime) -> i128 {
81 dt.to_zoned(TimeZone::UTC)
82 .expect("DateTime bound out of jiff::Timestamp range")
83 .timestamp()
84 .as_nanosecond()
85}
86
87fn nanos_to_datetime(n: i128) -> DateTime {
89 Timestamp::from_nanosecond(n)
90 .unwrap()
91 .to_zoned(TimeZone::UTC)
92 .datetime()
93}
94
95pub struct DateTimeGenerator {
97 min_value: DateTime,
98 max_value: DateTime,
99}
100
101impl DateTimeGenerator {
102 pub fn min_value(mut self, min: DateTime) -> Self {
104 self.min_value = min;
105 self
106 }
107
108 pub fn max_value(mut self, max: DateTime) -> Self {
110 self.max_value = max;
111 self
112 }
113}
114
115impl Generator<DateTime> for DateTimeGenerator {
116 fn as_basic(&self) -> Option<BasicGenerator<'_, DateTime>> {
117 assert!(
118 self.min_value <= self.max_value,
119 "Cannot have max_value < min_value"
120 );
121 let schema = cbor_map! {
122 "type" => "integer",
123 "min_value" => datetime_to_nanos(self.min_value),
124 "max_value" => datetime_to_nanos(self.max_value),
125 };
126 Some(BasicGenerator::new(schema, |raw| {
127 let n: i128 = deserialize_value(raw);
128 nanos_to_datetime(n)
129 }))
130 }
131}
132
133pub fn datetimes() -> DateTimeGenerator {
151 DateTimeGenerator {
152 min_value: Timestamp::MIN.to_zoned(TimeZone::UTC).datetime(),
153 max_value: Timestamp::MAX.to_zoned(TimeZone::UTC).datetime(),
154 }
155}
156
157pub struct TimestampGenerator {
159 min_value: Timestamp,
160 max_value: Timestamp,
161}
162
163impl TimestampGenerator {
164 pub fn min_value(mut self, min: Timestamp) -> Self {
166 self.min_value = min;
167 self
168 }
169
170 pub fn max_value(mut self, max: Timestamp) -> Self {
172 self.max_value = max;
173 self
174 }
175
176 fn build_schema(&self) -> Value {
177 assert!(
178 self.min_value <= self.max_value,
179 "Cannot have max_value < min_value"
180 );
181 cbor_map! {
182 "type" => "integer",
183 "min_value" => self.min_value.as_nanosecond(),
184 "max_value" => self.max_value.as_nanosecond(),
185 }
186 }
187}
188
189impl Generator<Timestamp> for TimestampGenerator {
190 fn as_basic(&self) -> Option<BasicGenerator<'_, Timestamp>> {
191 Some(BasicGenerator::new(self.build_schema(), |raw| {
192 let nanos: i128 = deserialize_value(raw);
193 Timestamp::from_nanosecond(nanos).unwrap()
194 }))
195 }
196}
197
198pub fn timestamps() -> TimestampGenerator {
216 TimestampGenerator {
217 min_value: Timestamp::MIN,
218 max_value: Timestamp::MAX,
219 }
220}
221
222const SPAN_NANOS_LIMIT: i64 = i64::MAX;
224
225pub struct SpanGenerator {
227 min_nanos: i64,
228 max_nanos: i64,
229}
230
231impl SpanGenerator {
232 pub fn min_nanoseconds(mut self, min: i64) -> Self {
234 self.min_nanos = min;
235 self
236 }
237
238 pub fn max_nanoseconds(mut self, max: i64) -> Self {
240 self.max_nanos = max;
241 self
242 }
243
244 fn build_schema(&self) -> Value {
245 assert!(
246 self.min_nanos <= self.max_nanos,
247 "Cannot have max_nanoseconds < min_nanoseconds"
248 );
249 assert!(
250 self.min_nanos >= -SPAN_NANOS_LIMIT,
251 "min_nanoseconds must be >= -i64::MAX (the Span nanosecond limit)"
252 );
253 cbor_map! {
254 "type" => "integer",
255 "min_value" => self.min_nanos,
256 "max_value" => self.max_nanos,
257 }
258 }
259}
260
261impl Generator<Span> for SpanGenerator {
262 fn as_basic(&self) -> Option<BasicGenerator<'_, Span>> {
263 Some(BasicGenerator::new(self.build_schema(), |raw| {
264 let nanos: i64 = deserialize_value(raw);
265 Span::new().try_nanoseconds(nanos).unwrap()
266 }))
267 }
268}
269
270pub fn spans() -> SpanGenerator {
286 SpanGenerator {
287 min_nanos: -SPAN_NANOS_LIMIT,
288 max_nanos: SPAN_NANOS_LIMIT,
289 }
290}
291
292fn nanos_to_signed_duration(n: i128) -> SignedDuration {
294 let secs = (n / 1_000_000_000) as i64;
295 let subnanos = (n % 1_000_000_000) as i32;
296 SignedDuration::new(secs, subnanos)
297}
298
299pub struct SignedDurationGenerator {
301 min_value: SignedDuration,
302 max_value: SignedDuration,
303}
304
305impl SignedDurationGenerator {
306 pub fn min_value(mut self, min: SignedDuration) -> Self {
308 self.min_value = min;
309 self
310 }
311
312 pub fn max_value(mut self, max: SignedDuration) -> Self {
314 self.max_value = max;
315 self
316 }
317
318 fn build_schema(&self) -> Value {
319 assert!(
320 self.min_value <= self.max_value,
321 "Cannot have max_value < min_value"
322 );
323 cbor_map! {
324 "type" => "integer",
325 "min_value" => self.min_value.as_nanos(),
326 "max_value" => self.max_value.as_nanos(),
327 }
328 }
329}
330
331impl Generator<SignedDuration> for SignedDurationGenerator {
332 fn as_basic(&self) -> Option<BasicGenerator<'_, SignedDuration>> {
333 Some(BasicGenerator::new(self.build_schema(), |raw| {
334 let nanos: i128 = deserialize_value(raw);
335 nanos_to_signed_duration(nanos)
336 }))
337 }
338}
339
340pub fn signed_durations() -> SignedDurationGenerator {
355 SignedDurationGenerator {
356 min_value: SignedDuration::MIN,
357 max_value: SignedDuration::MAX,
358 }
359}
360
361const OFFSET_MIN_SECS: i32 = -93_599;
363const OFFSET_MAX_SECS: i32 = 93_599;
365
366pub struct OffsetGenerator {
368 min_value: Offset,
369 max_value: Offset,
370}
371
372impl OffsetGenerator {
373 pub fn min_value(mut self, min: Offset) -> Self {
375 self.min_value = min;
376 self
377 }
378
379 pub fn max_value(mut self, max: Offset) -> Self {
381 self.max_value = max;
382 self
383 }
384
385 fn build_schema(&self) -> Value {
386 assert!(
387 self.min_value.seconds() <= self.max_value.seconds(),
388 "Cannot have max_value < min_value"
389 );
390 cbor_map! {
391 "type" => "integer",
392 "min_value" => self.min_value.seconds(),
393 "max_value" => self.max_value.seconds(),
394 }
395 }
396}
397
398impl Generator<Offset> for OffsetGenerator {
399 fn as_basic(&self) -> Option<BasicGenerator<'_, Offset>> {
400 Some(BasicGenerator::new(self.build_schema(), |raw| {
401 let secs: i32 = deserialize_value(raw);
402 Offset::from_seconds(secs).unwrap()
403 }))
404 }
405}
406
407pub fn offsets() -> OffsetGenerator {
425 OffsetGenerator {
426 min_value: Offset::from_seconds(OFFSET_MIN_SECS).unwrap(),
427 max_value: Offset::from_seconds(OFFSET_MAX_SECS).unwrap(),
428 }
429}
430
431pub struct ZonedGenerator<TS = TimestampGenerator, TZ = BoxedGenerator<'static, TimeZone>> {
433 timestamp_gen: TS,
434 timezone_gen: TZ,
435}
436
437impl<TS, TZ> ZonedGenerator<TS, TZ> {
438 pub fn timestamps<TS2>(self, timestamp_gen: TS2) -> ZonedGenerator<TS2, TZ>
440 where
441 TS2: Generator<Timestamp>,
442 {
443 ZonedGenerator {
444 timestamp_gen,
445 timezone_gen: self.timezone_gen,
446 }
447 }
448
449 pub fn timezones<TZ2>(self, timezone_gen: TZ2) -> ZonedGenerator<TS, TZ2>
451 where
452 TZ2: Generator<TimeZone>,
453 {
454 ZonedGenerator {
455 timestamp_gen: self.timestamp_gen,
456 timezone_gen,
457 }
458 }
459}
460
461impl<TS, TZ> Generator<Zoned> for ZonedGenerator<TS, TZ>
462where
463 TS: Generator<Timestamp>,
464 TZ: Generator<TimeZone>,
465{
466 fn as_basic(&self) -> Option<BasicGenerator<'_, Zoned>> {
467 let ts_basic = self.timestamp_gen.as_basic()?;
468 let tz_basic = self.timezone_gen.as_basic()?;
469 let schema = cbor_map! {
470 "type" => "tuple",
471 "elements" => cbor_array![
472 ts_basic.schema().clone(),
473 tz_basic.schema().clone(),
474 ],
475 };
476 Some(BasicGenerator::new(schema, move |raw| {
477 let [ts_raw, tz_raw]: [Value; 2] = raw.into_array().unwrap().try_into().unwrap();
478 let ts = ts_basic.parse_raw(ts_raw);
479 let tz = tz_basic.parse_raw(tz_raw);
480 Zoned::new(ts, tz)
481 }))
482 }
483}
484
485pub fn zoneds() -> ZonedGenerator {
501 use crate::generators::DefaultGenerator;
502 ZonedGenerator {
503 timestamp_gen: timestamps(),
504 timezone_gen: <TimeZone as DefaultGenerator>::default_generator(),
505 }
506}