drizzle_seed/
generators.rs

1//! Ready-made generators to pass to `SeedConfig::generator`.
2//!
3//! Every function returns a [`Generator`]. The [`GeneratorExt`] methods
4//! combine them: [`nullable`](GeneratorExt::nullable) mixes in `NULL`s and
5//! [`map`](GeneratorExt::map) post-processes values.
6//!
7//! # Examples
8//!
9//! ```rust
10//! use drizzle_seed::generators::{self, GeneratorExt};
11//! use drizzle_seed::{Generator, SeedValue};
12//! # use rand::SeedableRng;
13//! # let mut rng = rand::rngs::StdRng::seed_from_u64(1);
14//!
15//! let age = generators::int(18..=90);
16//! let plan = generators::one_of(["free", "pro", "team"]);
17//! let bio = generators::words(5..=20).nullable(0.3);
18//! let slug = generators::from_fn(|_rng, index| format!("post-{}", index + 1));
19//!
20//! assert!(matches!(age.generate(&mut rng, 0, "INTEGER"), SeedValue::Integer(18..=90)));
21//! assert!(matches!(plan.generate(&mut rng, 0, "TEXT"), SeedValue::Text(_)));
22//! assert!(matches!(bio.generate(&mut rng, 0, "TEXT"), SeedValue::Text(_) | SeedValue::Null));
23//! assert_eq!(slug.generate(&mut rng, 4, "TEXT"), SeedValue::Text("post-5".into()));
24//! ```
25//!
26//! With a schema, set them per column:
27//!
28//! ```text
29//! SeedConfig::postgres(&schema)
30//!     .generator(&schema.users.age, generators::int(18..=90))
31//!     .generator(&schema.users.plan, generators::weighted([("free", 8), ("pro", 2)]))
32//!     .generator(&schema.users.nickname, generators::first_name().nullable(0.5))
33//!     .generate();
34//! ```
35
36use crate::generator::{Generator, GeneratorKind, RngCore, SeedValue, string};
37use core::range::RangeInclusive;
38use rand::Rng;
39
40/// A range accepted by the generator constructors: `a..b` or `a..=b`,
41/// written with range syntax or as a `core::range` value.
42///
43/// Generators store it as a [`core::range::RangeInclusive`], which is
44/// `Copy` and has public `start` and `last` fields.
45pub trait IntoRange<T>: sealed::Sealed<T> {
46    /// The same values as an inclusive range, or `None` when it is empty.
47    fn into_range(self) -> Option<RangeInclusive<T>>;
48}
49
50mod sealed {
51    pub trait Sealed<T> {}
52}
53
54macro_rules! integer_ranges {
55    ($($ty:ty),*) => {$(
56        impl sealed::Sealed<$ty> for core::ops::Range<$ty> {}
57        impl IntoRange<$ty> for core::ops::Range<$ty> {
58            fn into_range(self) -> Option<RangeInclusive<$ty>> {
59                (self.start < self.end).then(|| RangeInclusive {
60                    start: self.start,
61                    last: self.end - 1,
62                })
63            }
64        }
65        impl sealed::Sealed<$ty> for core::range::Range<$ty> {}
66        impl IntoRange<$ty> for core::range::Range<$ty> {
67            fn into_range(self) -> Option<RangeInclusive<$ty>> {
68                core::ops::Range::from(self).into_range()
69            }
70        }
71        impl sealed::Sealed<$ty> for core::ops::RangeInclusive<$ty> {}
72        impl IntoRange<$ty> for core::ops::RangeInclusive<$ty> {
73            fn into_range(self) -> Option<RangeInclusive<$ty>> {
74                (!self.is_empty()).then(|| self.into())
75            }
76        }
77        impl sealed::Sealed<$ty> for RangeInclusive<$ty> {}
78        impl IntoRange<$ty> for RangeInclusive<$ty> {
79            fn into_range(self) -> Option<RangeInclusive<$ty>> {
80                (self.start <= self.last).then_some(self)
81            }
82        }
83    )*};
84}
85
86integer_ranges!(i64, usize);
87
88impl sealed::Sealed<f64> for core::ops::RangeInclusive<f64> {}
89impl IntoRange<f64> for core::ops::RangeInclusive<f64> {
90    fn into_range(self) -> Option<RangeInclusive<f64>> {
91        RangeInclusive::from(self).into_range()
92    }
93}
94impl sealed::Sealed<f64> for RangeInclusive<f64> {}
95impl IntoRange<f64> for RangeInclusive<f64> {
96    fn into_range(self) -> Option<RangeInclusive<f64>> {
97        (self.start.is_finite() && self.last.is_finite() && self.start <= self.last).then_some(self)
98    }
99}
100
101/// Combinators available on every [`Generator`].
102pub trait GeneratorExt: Generator + Sized {
103    /// Returns `NULL` with the given probability (0.0 to 1.0), otherwise a
104    /// value from `self`.
105    ///
106    /// # Panics
107    ///
108    /// Panics if `probability` is not between 0.0 and 1.0.
109    fn nullable(self, probability: f64) -> Nullable<Self> {
110        assert!(
111            (0.0..=1.0).contains(&probability),
112            "nullable probability must be between 0.0 and 1.0, got {probability}"
113        );
114        Nullable {
115            inner: self,
116            probability,
117        }
118    }
119
120    /// Transforms each generated value with `f`.
121    fn map<F, V>(self, f: F) -> Map<Self, F>
122    where
123        F: Fn(SeedValue) -> V + Send + Sync,
124        V: Into<SeedValue>,
125    {
126        Map { inner: self, f }
127    }
128}
129
130impl<G: Generator> GeneratorExt for G {}
131
132/// Random integers from a range. Created by [`int`].
133#[derive(Debug, Clone, Copy)]
134pub struct Int {
135    range: RangeInclusive<i64>,
136}
137
138/// Random integers in `range`, for example `int(18..=90)` or `int(0..10)`.
139///
140/// # Panics
141///
142/// Panics if the range is empty.
143#[must_use]
144pub fn int(range: impl IntoRange<i64>) -> Int {
145    Int {
146        range: range.into_range().expect("int range is empty"),
147    }
148}
149
150impl Generator for Int {
151    fn generate(&self, rng: &mut dyn RngCore, _index: usize, _sql_type: &str) -> SeedValue {
152        SeedValue::Integer(rng.random_range(self.range.start..=self.range.last))
153    }
154
155    fn name(&self) -> &'static str {
156        "int"
157    }
158}
159
160/// Random floating-point numbers from a range. Created by [`float`].
161#[derive(Debug, Clone, Copy)]
162pub struct Float {
163    range: RangeInclusive<f64>,
164    decimals: Option<u32>,
165}
166
167/// Random numbers in `range`, for example `float(0.0..=100.0)`.
168///
169/// # Panics
170///
171/// Panics if the range is empty or not finite.
172#[must_use]
173pub fn float(range: impl IntoRange<f64>) -> Float {
174    Float {
175        range: range
176            .into_range()
177            .expect("float range must be finite and non-empty"),
178        decimals: None,
179    }
180}
181
182impl Float {
183    /// Rounds each value to `decimals` digits after the point, for example
184    /// prices with `.decimals(2)`.
185    #[must_use]
186    pub const fn decimals(mut self, decimals: u32) -> Self {
187        self.decimals = Some(decimals);
188        self
189    }
190}
191
192impl Generator for Float {
193    fn generate(&self, rng: &mut dyn RngCore, _index: usize, _sql_type: &str) -> SeedValue {
194        let value = rng.random_range(self.range.start..=self.range.last);
195        let value = match self.decimals {
196            Some(decimals) => {
197                let factor = 10f64.powi(i32::try_from(decimals).unwrap_or(i32::MAX));
198                ((value * factor).round() / factor).clamp(self.range.start, self.range.last)
199            }
200            None => value,
201        };
202        SeedValue::Float(value)
203    }
204
205    fn name(&self) -> &'static str {
206        "float"
207    }
208}
209
210/// Random booleans. Created by [`boolean`].
211#[derive(Debug, Clone, Copy)]
212pub struct Boolean {
213    probability_true: f64,
214}
215
216/// `true` with the given probability (0.0 to 1.0), otherwise `false`.
217///
218/// # Panics
219///
220/// Panics if `probability_true` is not between 0.0 and 1.0.
221#[must_use]
222pub fn boolean(probability_true: f64) -> Boolean {
223    assert!(
224        (0.0..=1.0).contains(&probability_true),
225        "boolean probability must be between 0.0 and 1.0, got {probability_true}"
226    );
227    Boolean { probability_true }
228}
229
230impl Generator for Boolean {
231    fn generate(&self, rng: &mut dyn RngCore, _index: usize, _sql_type: &str) -> SeedValue {
232        SeedValue::Bool(rng.random_bool(self.probability_true))
233    }
234
235    fn name(&self) -> &'static str {
236        "boolean"
237    }
238}
239
240/// Picks one of a fixed set of values. Created by [`one_of`].
241#[derive(Debug, Clone)]
242pub struct OneOf {
243    values: Vec<SeedValue>,
244}
245
246/// Picks one of `values`, each equally likely, for example
247/// `one_of(["free", "pro"])`.
248///
249/// # Panics
250///
251/// Panics if `values` is empty.
252#[must_use]
253pub fn one_of<T: Into<SeedValue>>(values: impl IntoIterator<Item = T>) -> OneOf {
254    let values: Vec<SeedValue> = values.into_iter().map(Into::into).collect();
255    assert!(!values.is_empty(), "one_of needs at least one value");
256    OneOf { values }
257}
258
259impl Generator for OneOf {
260    fn generate(&self, rng: &mut dyn RngCore, _index: usize, _sql_type: &str) -> SeedValue {
261        self.values[rng.random_range(0..self.values.len())].clone()
262    }
263
264    fn name(&self) -> &'static str {
265        "one_of"
266    }
267}
268
269/// Picks values in proportion to their weights. Created by [`weighted`].
270#[derive(Debug, Clone)]
271pub struct Weighted {
272    values: Vec<SeedValue>,
273    cumulative: Vec<u64>,
274}
275
276/// Picks one of `values` in proportion to its weight, for example
277/// `weighted([("free", 8), ("pro", 2)])` gives `"free"` 80% of the time.
278///
279/// # Panics
280///
281/// Panics if `values` is empty or every weight is zero.
282#[must_use]
283pub fn weighted<T: Into<SeedValue>>(values: impl IntoIterator<Item = (T, u32)>) -> Weighted {
284    let mut total = 0u64;
285    let (values, cumulative) = values
286        .into_iter()
287        .map(|(value, weight)| {
288            total += u64::from(weight);
289            (value.into(), total)
290        })
291        .unzip();
292    assert!(total > 0, "weighted needs at least one non-zero weight");
293    Weighted { values, cumulative }
294}
295
296impl Generator for Weighted {
297    fn generate(&self, rng: &mut dyn RngCore, _index: usize, _sql_type: &str) -> SeedValue {
298        let total = *self.cumulative.last().unwrap_or(&1);
299        let pick = rng.random_range(0..total);
300        let index = self.cumulative.partition_point(|&end| end <= pick);
301        self.values[index].clone()
302    }
303
304    fn name(&self) -> &'static str {
305        "weighted"
306    }
307}
308
309/// The same value for every row. Created by [`constant`].
310#[derive(Debug, Clone)]
311pub struct Constant {
312    value: SeedValue,
313}
314
315/// The same `value` for every row.
316#[must_use]
317pub fn constant(value: impl Into<SeedValue>) -> Constant {
318    Constant {
319        value: value.into(),
320    }
321}
322
323impl Generator for Constant {
324    fn generate(&self, _rng: &mut dyn RngCore, _index: usize, _sql_type: &str) -> SeedValue {
325        self.value.clone()
326    }
327
328    fn name(&self) -> &'static str {
329        "constant"
330    }
331}
332
333/// Counts up from a start value. Created by [`sequence`].
334#[derive(Debug, Clone, Copy)]
335pub struct Sequence {
336    start: i64,
337    step: i64,
338}
339
340/// `start`, `start + 1`, ... by row; change the increment with
341/// [`step`](Sequence::step).
342#[must_use]
343pub const fn sequence(start: i64) -> Sequence {
344    Sequence { start, step: 1 }
345}
346
347impl Sequence {
348    /// Sets the increment between rows.
349    #[must_use]
350    pub const fn step(mut self, step: i64) -> Self {
351        self.step = step;
352        self
353    }
354}
355
356impl Generator for Sequence {
357    fn generate(&self, _rng: &mut dyn RngCore, index: usize, _sql_type: &str) -> SeedValue {
358        let index = i64::try_from(index).unwrap_or(i64::MAX);
359        SeedValue::Integer(self.start.saturating_add(index.saturating_mul(self.step)))
360    }
361
362    fn name(&self) -> &'static str {
363        "sequence"
364    }
365}
366
367/// Random lowercase letters. Created by [`text`].
368#[derive(Debug, Clone, Copy)]
369pub struct Text {
370    length: RangeInclusive<usize>,
371}
372
373/// A string of random lowercase letters whose length is in `length`, for
374/// example `text(5..=12)`.
375///
376/// # Panics
377///
378/// Panics if the range is empty.
379#[must_use]
380pub fn text(length: impl IntoRange<usize>) -> Text {
381    Text {
382        length: length.into_range().expect("text length range is empty"),
383    }
384}
385
386impl Generator for Text {
387    fn generate(&self, rng: &mut dyn RngCore, index: usize, sql_type: &str) -> SeedValue {
388        string::TextGen {
389            min_len: self.length.start,
390            max_len: self.length.last,
391        }
392        .generate(rng, index, sql_type)
393    }
394
395    fn name(&self) -> &'static str {
396        "text"
397    }
398}
399
400/// Lorem ipsum words. Created by [`words`].
401#[derive(Debug, Clone, Copy)]
402pub struct Words {
403    count: RangeInclusive<usize>,
404}
405
406/// Lorem ipsum: a number of words in `count`, separated by spaces, for
407/// example `words(5..=20)`.
408///
409/// # Panics
410///
411/// Panics if the range is empty.
412#[must_use]
413pub fn words(count: impl IntoRange<usize>) -> Words {
414    Words {
415        count: count.into_range().expect("words count range is empty"),
416    }
417}
418
419impl Generator for Words {
420    fn generate(&self, rng: &mut dyn RngCore, index: usize, sql_type: &str) -> SeedValue {
421        let words = rng.random_range(self.count.start..=self.count.last);
422        string::LoremGen { words }.generate(rng, index, sql_type)
423    }
424
425    fn name(&self) -> &'static str {
426        "words"
427    }
428}
429
430/// A generator from a closure. Created by [`from_fn`].
431pub struct FromFn<F> {
432    f: F,
433}
434
435/// Uses a closure as a generator. It gets the RNG and the 0-based row
436/// index, and returns anything that converts into a [`SeedValue`].
437///
438/// Draw randomness only from the given RNG, so the same seed gives the same
439/// rows. Its methods come from [`Rng`], re-exported by this
440/// crate.
441#[must_use]
442pub fn from_fn<F, V>(f: F) -> FromFn<F>
443where
444    F: Fn(&mut dyn RngCore, usize) -> V + Send + Sync,
445    V: Into<SeedValue>,
446{
447    FromFn { f }
448}
449
450impl<F, V> Generator for FromFn<F>
451where
452    F: Fn(&mut dyn RngCore, usize) -> V + Send + Sync,
453    V: Into<SeedValue>,
454{
455    fn generate(&self, rng: &mut dyn RngCore, index: usize, _sql_type: &str) -> SeedValue {
456        (self.f)(rng, index).into()
457    }
458
459    fn name(&self) -> &'static str {
460        "from_fn"
461    }
462}
463
464/// Mixes `NULL` into another generator. Created by
465/// [`GeneratorExt::nullable`].
466pub struct Nullable<G> {
467    inner: G,
468    probability: f64,
469}
470
471impl<G: Generator> Generator for Nullable<G> {
472    fn generate(&self, rng: &mut dyn RngCore, index: usize, sql_type: &str) -> SeedValue {
473        if rng.random_bool(self.probability) {
474            SeedValue::Null
475        } else {
476            self.inner.generate(rng, index, sql_type)
477        }
478    }
479
480    fn name(&self) -> &'static str {
481        self.inner.name()
482    }
483}
484
485/// Transforms another generator's values. Created by [`GeneratorExt::map`].
486pub struct Map<G, F> {
487    inner: G,
488    f: F,
489}
490
491impl<G, F, V> Generator for Map<G, F>
492where
493    G: Generator,
494    F: Fn(SeedValue) -> V + Send + Sync,
495    V: Into<SeedValue>,
496{
497    fn generate(&self, rng: &mut dyn RngCore, index: usize, sql_type: &str) -> SeedValue {
498        (self.f)(self.inner.generate(rng, index, sql_type)).into()
499    }
500
501    fn name(&self) -> &'static str {
502        self.inner.name()
503    }
504}
505
506macro_rules! kind_constructors {
507    ($($(#[$doc:meta])* $name:ident => $kind:ident,)*) => {$(
508        $(#[$doc])*
509        #[must_use]
510        pub const fn $name() -> GeneratorKind {
511            GeneratorKind::$kind
512        }
513    )*};
514}
515
516kind_constructors! {
517    /// `first.last{row}@example.com` (or another reserved example domain);
518    /// the row number keeps emails unique.
519    email => Email,
520    /// A first name.
521    first_name => FirstName,
522    /// A last name.
523    last_name => LastName,
524    /// A first and last name.
525    full_name => FullName,
526    /// A US-style phone number, `(555) 555-5555`.
527    phone => Phone,
528    /// A city name.
529    city => City,
530    /// A country name.
531    country => Country,
532    /// A street address, `123 Name Street`.
533    address => Address,
534    /// A company name, such as `Smith Inc`.
535    company => Company,
536    /// A job title.
537    job_title => JobTitle,
538    /// A short sentence-case phrase, such as a post title.
539    title => Title,
540    /// An `https://` URL on a reserved example domain, unique per row.
541    url => Url,
542    /// A lowercase username, `first_last{row}`, unique per row.
543    username => Username,
544    /// A URL slug, `dolor-sit-amet-{row}`, unique per row.
545    slug => Slug,
546    /// A random UUID v4 as text.
547    uuid => Uuid,
548    /// A small JSON object as text.
549    json => Json,
550    /// A date between 2000 and 2030 as `YYYY-MM-DD` (Unix milliseconds for
551    /// an integer column).
552    date => Date,
553    /// A timestamp between 2000 and 2030 as `YYYY-MM-DD HH:MM:SS` (Unix
554    /// milliseconds for an integer column).
555    timestamp => Timestamp,
556    /// A time of day, `HH:MM:SS`.
557    time => Time,
558}
559
560#[cfg(test)]
561mod tests {
562    use super::*;
563    use rand::SeedableRng;
564    use rand::rngs::StdRng;
565
566    fn values(generator: &impl Generator, rows: usize) -> Vec<SeedValue> {
567        let mut rng = StdRng::seed_from_u64(7);
568        (0..rows)
569            .map(|index| generator.generate(&mut rng, index, "TEXT"))
570            .collect()
571    }
572
573    #[test]
574    fn ranges_are_respected() {
575        assert!(
576            values(&int(-3..=3), 200)
577                .iter()
578                .all(|value| matches!(value, SeedValue::Integer(-3..=3)))
579        );
580        assert!(
581            values(&float(1.0..=2.0).decimals(1), 200)
582                .iter()
583                .all(|value| {
584                    matches!(value, SeedValue::Float(number) if (1.0..=2.0).contains(number)
585                && (number * 10.0).fract() == 0.0)
586                })
587        );
588        assert!(values(&text(2..=4), 200).iter().all(|value| {
589            matches!(value, SeedValue::Text(text) if (2..=4).contains(&text.len()))
590        }));
591    }
592
593    #[test]
594    fn exclusive_inclusive_and_core_range_forms_agree() {
595        let exclusive = values(&int(0..3), 200);
596        assert_eq!(exclusive, values(&int(0..=2), 200));
597        assert_eq!(
598            exclusive,
599            values(&int(core::range::Range { start: 0, end: 3 }), 200)
600        );
601        assert_eq!(
602            exclusive,
603            values(&int(core::range::RangeInclusive { start: 0, last: 2 }), 200)
604        );
605        assert_eq!(values(&text(3..4), 20), values(&text(3..=3), 20));
606    }
607
608    #[test]
609    #[should_panic(expected = "int range is empty")]
610    fn empty_ranges_are_rejected() {
611        let _ = int(5..5);
612    }
613
614    #[test]
615    fn choices_and_weights() {
616        let picked = values(&one_of(["a", "b"]), 200);
617        assert!(picked.contains(&SeedValue::from("a")) && picked.contains(&SeedValue::from("b")));
618
619        let picked = values(&weighted([("never", 0), ("always", 3)]), 100);
620        assert!(
621            picked
622                .iter()
623                .all(|value| *value == SeedValue::from("always"))
624        );
625    }
626
627    #[test]
628    fn sequence_constant_and_closures() {
629        assert_eq!(
630            values(&sequence(10).step(5), 3),
631            [10, 15, 20].map(SeedValue::from)
632        );
633        assert_eq!(
634            values(&constant(true), 2),
635            [true, true].map(SeedValue::from)
636        );
637        assert_eq!(
638            values(&from_fn(|_rng, index| index as i64 * 2), 3),
639            [0, 2, 4].map(SeedValue::from)
640        );
641    }
642
643    #[test]
644    fn combinators() {
645        let mixed = values(&int(1..=1).nullable(0.5), 200);
646        assert!(mixed.contains(&SeedValue::Null) && mixed.contains(&SeedValue::Integer(1)));
647        assert!(
648            values(&int(1..=1).nullable(0.0), 50)
649                .iter()
650                .all(|value| *value == SeedValue::Integer(1))
651        );
652        let mapped = values(
653            &first_name().map(|value| match value {
654                SeedValue::Text(name) => name.to_uppercase(),
655                _ => unreachable!(),
656            }),
657            5,
658        );
659        assert!(mapped.iter().all(|value| matches!(
660            value,
661            SeedValue::Text(name) if name.chars().all(|c| !c.is_lowercase())
662        )));
663    }
664}
drizzle_seed/generators.rs

drizzle_seed/
generators.rs
