hegel/generators/
strings.rs

1use super::{BasicGenerator, Generator, TestCase};
2use crate::cbor_utils::{cbor_array, cbor_map, map_extend, map_insert};
3use ciborium::Value;
4
5/// Categories that include surrogate codepoints. Rust strings cannot contain
6/// surrogates, so these are forbidden in `categories()`.
7const SURROGATE_CATEGORIES: &[&str] = &["Cs", "C"];
8
9/// Shared character filtering fields used by both [`TextGenerator`] and
10/// [`CharactersGenerator`].
11struct CharacterFields {
12    codec: Option<String>,
13    min_codepoint: Option<u32>,
14    max_codepoint: Option<u32>,
15    categories: Option<Vec<String>>,
16    exclude_categories: Option<Vec<String>>,
17    include_characters: Option<String>,
18    exclude_characters: Option<String>,
19}
20
21impl CharacterFields {
22    fn new() -> Self {
23        CharacterFields {
24            codec: None,
25            min_codepoint: None,
26            max_codepoint: None,
27            categories: None,
28            exclude_categories: None,
29            include_characters: None,
30            exclude_characters: None,
31        }
32    }
33
34    /// Build a schema map containing the character filtering fields.
35    fn to_schema(&self) -> Value {
36        let mut schema = cbor_map! {};
37        if let Some(ref codec) = self.codec {
38            map_insert(&mut schema, "codec", codec.as_str());
39        }
40        if let Some(min_cp) = self.min_codepoint {
41            map_insert(&mut schema, "min_codepoint", min_cp as u64);
42        }
43        if let Some(max_cp) = self.max_codepoint {
44            map_insert(&mut schema, "max_codepoint", max_cp as u64);
45        }
46        if let Some(ref cats) = self.categories {
47            for cat in cats {
48                assert!(
49                    !SURROGATE_CATEGORIES.contains(&cat.as_str()),
50                    "Category \"{cat}\" includes surrogate codepoints (Cs), \
51                     which Rust strings cannot represent."
52                );
53            }
54            let arr = Value::Array(cats.iter().map(|c| Value::from(c.as_str())).collect());
55            map_insert(&mut schema, "categories", arr);
56        } else {
57            // Always exclude surrogates (Cs) since Rust strings cannot contain them.
58            let mut excl = self.exclude_categories.clone().unwrap_or_default();
59            if !excl.iter().any(|c| c == "Cs") {
60                excl.push("Cs".to_string());
61            }
62            let arr = Value::Array(excl.iter().map(|c| Value::from(c.as_str())).collect());
63            map_insert(&mut schema, "exclude_categories", arr);
64        }
65        if let Some(ref incl) = self.include_characters {
66            map_insert(&mut schema, "include_characters", incl.as_str());
67        }
68        if let Some(ref excl) = self.exclude_characters {
69            map_insert(&mut schema, "exclude_characters", excl.as_str());
70        }
71        schema
72    }
73}
74
75/// Generator for Unicode text strings. Created by [`text()`].
76pub struct TextGenerator {
77    min_size: usize,
78    max_size: Option<usize>,
79    char_fields: CharacterFields,
80    alphabet_called: bool,
81    char_param_called: bool,
82}
83
84impl TextGenerator {
85    /// Set the minimum length in characters.
86    pub fn min_size(mut self, min_size: usize) -> Self {
87        self.min_size = min_size;
88        self
89    }
90
91    /// Set the maximum length in characters.
92    pub fn max_size(mut self, max_size: usize) -> Self {
93        self.max_size = Some(max_size);
94        self
95    }
96
97    /// Use a fixed set of characters. Each character in the generated string
98    /// will be a member of the alphabet.
99    ///
100    /// Mutually exclusive with the character filtering methods like `codec`,
101    /// `categories`, `min_codepoint`, etc.
102    pub fn alphabet(mut self, chars: &str) -> Self {
103        self.char_fields = CharacterFields {
104            codec: None,
105            min_codepoint: None,
106            max_codepoint: None,
107            categories: Some(vec![]),
108            exclude_categories: None,
109            include_characters: Some(chars.to_string()),
110            exclude_characters: None,
111        };
112        self.alphabet_called = true;
113        self
114    }
115
116    /// Restrict to characters encodable in this codec (e.g. `"ascii"`, `"utf-8"`, `"latin-1"`).
117    pub fn codec(mut self, codec: &str) -> Self {
118        self.char_param_called = true;
119        self.char_fields.codec = Some(codec.to_string());
120        self
121    }
122
123    /// Set the minimum Unicode codepoint.
124    pub fn min_codepoint(mut self, min_codepoint: u32) -> Self {
125        self.char_param_called = true;
126        self.char_fields.min_codepoint = Some(min_codepoint);
127        self
128    }
129
130    /// Set the maximum Unicode codepoint.
131    pub fn max_codepoint(mut self, max_codepoint: u32) -> Self {
132        self.char_param_called = true;
133        self.char_fields.max_codepoint = Some(max_codepoint);
134        self
135    }
136
137    /// Include only characters from these Unicode general categories (e.g. `["L", "Nd"]`).
138    ///
139    /// Mutually exclusive with [`exclude_categories`](Self::exclude_categories).
140    pub fn categories(mut self, categories: &[&str]) -> Self {
141        self.char_param_called = true;
142        self.char_fields.categories = Some(categories.iter().map(|s| s.to_string()).collect());
143        self
144    }
145
146    /// Exclude characters from these Unicode general categories.
147    ///
148    /// Mutually exclusive with [`categories`](Self::categories).
149    pub fn exclude_categories(mut self, exclude_categories: &[&str]) -> Self {
150        self.char_param_called = true;
151        self.char_fields.exclude_categories =
152            Some(exclude_categories.iter().map(|s| s.to_string()).collect());
153        self
154    }
155
156    /// Always include these specific characters, even if excluded by other filters.
157    pub fn include_characters(mut self, include_characters: &str) -> Self {
158        self.char_param_called = true;
159        self.char_fields.include_characters = Some(include_characters.to_string());
160        self
161    }
162
163    /// Always exclude these specific characters.
164    pub fn exclude_characters(mut self, exclude_characters: &str) -> Self {
165        self.char_param_called = true;
166        self.char_fields.exclude_characters = Some(exclude_characters.to_string());
167        self
168    }
169
170    fn build_schema(&self) -> Value {
171        assert!(
172            !(self.alphabet_called && self.char_param_called),
173            "Cannot combine .alphabet() with character methods."
174        );
175        if let Some(max) = self.max_size {
176            assert!(self.min_size <= max, "Cannot have max_size < min_size");
177        }
178
179        let mut schema = cbor_map! {
180            "type" => "string",
181            "min_size" => self.min_size as u64
182        };
183
184        if let Some(max) = self.max_size {
185            map_insert(&mut schema, "max_size", max as u64);
186        }
187        map_extend(&mut schema, self.char_fields.to_schema());
188
189        schema
190    }
191}
192
193impl Generator<String> for TextGenerator {
194    fn as_basic(&self) -> Option<BasicGenerator<'_, String>> {
195        Some(BasicGenerator::new(
196            self.build_schema(),
197            super::deserialize_value,
198        ))
199    }
200}
201
202/// Generate arbitrary Unicode text strings.
203///
204/// See [`TextGenerator`] for builder methods.
205pub fn text() -> TextGenerator {
206    TextGenerator {
207        min_size: 0,
208        max_size: None,
209        char_fields: CharacterFields::new(),
210        alphabet_called: false,
211        char_param_called: false,
212    }
213}
214
215/// Generator for single Unicode characters ([`char`]). Created by [`characters()`].
216pub struct CharactersGenerator {
217    char_fields: CharacterFields,
218}
219
220impl CharactersGenerator {
221    /// Restrict to characters encodable in this codec (e.g. `"ascii"`, `"utf-8"`, `"latin-1"`).
222    pub fn codec(mut self, codec: &str) -> Self {
223        self.char_fields.codec = Some(codec.to_string());
224        self
225    }
226
227    /// Set the minimum Unicode codepoint.
228    pub fn min_codepoint(mut self, min_codepoint: u32) -> Self {
229        self.char_fields.min_codepoint = Some(min_codepoint);
230        self
231    }
232
233    /// Set the maximum Unicode codepoint.
234    pub fn max_codepoint(mut self, max_codepoint: u32) -> Self {
235        self.char_fields.max_codepoint = Some(max_codepoint);
236        self
237    }
238
239    /// Include only characters from these Unicode general categories (e.g. `["L", "Nd"]`).
240    ///
241    /// Mutually exclusive with [`exclude_categories`](Self::exclude_categories).
242    pub fn categories(mut self, categories: &[&str]) -> Self {
243        self.char_fields.categories = Some(categories.iter().map(|s| s.to_string()).collect());
244        self
245    }
246
247    /// Exclude characters from these Unicode general categories.
248    ///
249    /// Mutually exclusive with [`categories`](Self::categories).
250    pub fn exclude_categories(mut self, exclude_categories: &[&str]) -> Self {
251        self.char_fields.exclude_categories =
252            Some(exclude_categories.iter().map(|s| s.to_string()).collect());
253        self
254    }
255
256    /// Always include these specific characters, even if excluded by other filters.
257    pub fn include_characters(mut self, include_characters: &str) -> Self {
258        self.char_fields.include_characters = Some(include_characters.to_string());
259        self
260    }
261
262    /// Always exclude these specific characters.
263    pub fn exclude_characters(mut self, exclude_characters: &str) -> Self {
264        self.char_fields.exclude_characters = Some(exclude_characters.to_string());
265        self
266    }
267
268    fn build_schema(&self) -> Value {
269        let mut schema = cbor_map! {
270            "type" => "string",
271            "min_size" => 1u64,
272            "max_size" => 1u64
273        };
274        map_extend(&mut schema, self.char_fields.to_schema());
275        schema
276    }
277
278    /// Build a standalone schema for use as a regex alphabet constraint.
279    pub(super) fn build_alphabet_schema(&self) -> Value {
280        self.char_fields.to_schema()
281    }
282}
283
284fn parse_char(raw: Value) -> char {
285    let s: String = super::deserialize_value(raw);
286    let mut chars = s.chars();
287    let c = chars
288        .next()
289        .expect("expected a single character, got empty string");
290    assert!(
291        chars.next().is_none(),
292        "expected a single character, got multiple"
293    );
294    c
295}
296
297impl Generator<char> for CharactersGenerator {
298    fn as_basic(&self) -> Option<BasicGenerator<'_, char>> {
299        Some(BasicGenerator::new(self.build_schema(), parse_char))
300    }
301}
302
303/// Generate single Unicode characters ([`char`]).
304///
305/// See [`CharactersGenerator`] for builder methods.
306pub fn characters() -> CharactersGenerator {
307    CharactersGenerator {
308        char_fields: CharacterFields::new(),
309    }
310}
311
312/// Generator for strings matching a regex pattern. Created by [`from_regex()`].
313///
314/// By default generates strings that contain a match. Use [`fullmatch()`](Self::fullmatch)
315/// to require the entire string to match.
316pub struct RegexGenerator {
317    pattern: String,
318    fullmatch: bool,
319    alphabet: Option<CharactersGenerator>,
320}
321
322impl RegexGenerator {
323    /// Set whether the entire string must match the pattern, not just contain a match.
324    pub fn fullmatch(mut self, fullmatch: bool) -> Self {
325        self.fullmatch = fullmatch;
326        self
327    }
328
329    /// Constrain which characters may appear in generated strings.
330    pub fn alphabet(mut self, alphabet: CharactersGenerator) -> Self {
331        self.alphabet = Some(alphabet);
332        self
333    }
334
335    // nocov start
336    fn build_schema(&self) -> Value {
337        let mut schema = cbor_map! {
338            "type" => "regex",
339            "pattern" => self.pattern.as_str(),
340            "fullmatch" => self.fullmatch
341        // nocov end
342        };
343
344        if let Some(ref alphabet) = self.alphabet {
345            map_insert(&mut schema, "alphabet", alphabet.build_alphabet_schema());
346        }
347
348        schema
349    }
350}
351
352impl Generator<String> for RegexGenerator {
353    // nocov start
354    fn as_basic(&self) -> Option<BasicGenerator<'_, String>> {
355        Some(BasicGenerator::new(
356            self.build_schema(),
357            super::deserialize_value,
358            // nocov end
359        ))
360    }
361}
362
363/// Generate strings matching a regex pattern.
364///
365/// See [`RegexGenerator`] for builder methods.
366// nocov start
367pub fn from_regex(pattern: &str) -> RegexGenerator {
368    RegexGenerator {
369        pattern: pattern.to_string(),
370        fullmatch: false,
371        alphabet: None,
372        // nocov end
373    }
374}
375
376/// Generator for arbitrary byte sequences. Created by [`binary()`].
377pub struct BinaryGenerator {
378    min_size: usize,
379    max_size: Option<usize>,
380}
381
382impl BinaryGenerator {
383    /// Set the minimum length in bytes.
384    pub fn min_size(mut self, min_size: usize) -> Self {
385        self.min_size = min_size;
386        self
387    }
388
389    /// Set the maximum length in bytes.
390    pub fn max_size(mut self, max_size: usize) -> Self {
391        self.max_size = Some(max_size);
392        self
393    }
394
395    fn build_schema(&self) -> Value {
396        if let Some(max) = self.max_size {
397            assert!(self.min_size <= max, "Cannot have max_size < min_size");
398        }
399
400        let mut schema = cbor_map! {
401            "type" => "binary",
402            "min_size" => self.min_size as u64
403        };
404
405        if let Some(max) = self.max_size {
406            map_insert(&mut schema, "max_size", max as u64);
407        }
408
409        schema
410    }
411}
412
413fn parse_binary(raw: Value) -> Vec<u8> {
414    match raw {
415        Value::Bytes(bytes) => bytes,
416        _ => panic!("expected Value::Bytes, got {:?}", raw), // nocov
417    }
418}
419
420impl Generator<Vec<u8>> for BinaryGenerator {
421    fn as_basic(&self) -> Option<BasicGenerator<'_, Vec<u8>>> {
422        Some(BasicGenerator::new(self.build_schema(), parse_binary))
423    }
424}
425
426/// Generate arbitrary byte sequences (`Vec<u8>`).
427///
428/// See [`BinaryGenerator`] for builder methods.
429pub fn binary() -> BinaryGenerator {
430    BinaryGenerator {
431        min_size: 0,
432        max_size: None,
433    }
434}
435
436/// Generator for email address strings. Created by [`emails()`].
437pub struct EmailGenerator;
438
439impl Generator<String> for EmailGenerator {
440    fn as_basic(&self) -> Option<BasicGenerator<'_, String>> {
441        Some(BasicGenerator::new(cbor_map! {"type" => "email"}, |raw| {
442            super::deserialize_value(raw)
443        }))
444    }
445}
446
447/// Generate email address strings.
448pub fn emails() -> EmailGenerator {
449    EmailGenerator
450}
451
452/// Generator for URL strings. Created by [`urls()`].
453pub struct UrlGenerator;
454
455impl Generator<String> for UrlGenerator {
456    fn as_basic(&self) -> Option<BasicGenerator<'_, String>> {
457        Some(BasicGenerator::new(cbor_map! {"type" => "url"}, |raw| {
458            super::deserialize_value(raw)
459        }))
460    }
461}
462
463/// Generate URL strings.
464pub fn urls() -> UrlGenerator {
465    UrlGenerator
466}
467
468/// Generator for domain name strings. Created by [`domains()`].
469pub struct DomainGenerator {
470    max_length: usize,
471}
472
473impl DomainGenerator {
474    /// Set the maximum length (must be between 4 and 255).
475    pub fn max_length(mut self, max_length: usize) -> Self {
476        self.max_length = max_length;
477        self
478    }
479
480    fn build_schema(&self) -> Value {
481        assert!(
482            self.max_length >= 4 && self.max_length <= 255,
483            "max_length must be between 4 and 255"
484        );
485
486        cbor_map! {
487            "type" => "domain",
488            "max_length" => self.max_length as u64
489        }
490    }
491}
492
493impl Generator<String> for DomainGenerator {
494    fn as_basic(&self) -> Option<BasicGenerator<'_, String>> {
495        Some(BasicGenerator::new(self.build_schema(), |raw| {
496            super::deserialize_value(raw)
497        }))
498    }
499}
500
501/// Generate domain name strings.
502///
503/// See [`DomainGenerator`] for builder methods.
504pub fn domains() -> DomainGenerator {
505    DomainGenerator { max_length: 255 }
506}
507
508#[derive(Clone, Copy)]
509pub enum IpVersion {
510    V4,
511    V6,
512}
513
514/// Generator for IP address strings. Created by [`ip_addresses()`].
515///
516/// By default generates both IPv4 and IPv6 addresses.
517pub struct IpAddressGenerator {
518    version: Option<IpVersion>,
519}
520
521impl IpAddressGenerator {
522    /// Only generate IPv4 addresses.
523    pub fn v4(mut self) -> Self {
524        self.version = Some(IpVersion::V4);
525        self
526    }
527
528    /// Only generate IPv6 addresses.
529    pub fn v6(mut self) -> Self {
530        self.version = Some(IpVersion::V6);
531        self
532    }
533
534    fn build_schema(&self) -> Value {
535        match self.version {
536            Some(IpVersion::V4) => cbor_map! {"type" => "ip_address", "version" => 4u64},
537            Some(IpVersion::V6) => cbor_map! {"type" => "ip_address", "version" => 6u64},
538            None => cbor_map! {
539                "type" => "one_of",
540                "generators" => cbor_array![
541                    cbor_map!{"type" => "ip_address", "version" => 4u64},
542                    cbor_map!{"type" => "ip_address", "version" => 6u64}
543                ]
544            },
545        }
546    }
547}
548
549impl Generator<String> for IpAddressGenerator {
550    fn as_basic(&self) -> Option<BasicGenerator<'_, String>> {
551        let version = self.version;
552        Some(BasicGenerator::new(self.build_schema(), move |raw| {
553            // one_of returns [index, value]; extract the value for the mixed case.
554            let inner = if version.is_none() {
555                raw.into_array().unwrap().into_iter().nth(1).unwrap()
556            } else {
557                raw
558            };
559            super::deserialize_value(inner)
560        }))
561    }
562}
563
564/// Generate IP address strings (IPv4 or IPv6).
565///
566/// See [`IpAddressGenerator`] for builder methods.
567pub fn ip_addresses() -> IpAddressGenerator {
568    IpAddressGenerator { version: None }
569}
570
571/// Generator for date strings in YYYY-MM-DD format. Created by [`dates()`].
572pub struct DateGenerator;
573
574impl Generator<String> for DateGenerator {
575    // nocov start
576    fn as_basic(&self) -> Option<BasicGenerator<'_, String>> {
577        Some(BasicGenerator::new(cbor_map! {"type" => "date"}, |raw| {
578            super::deserialize_value(raw)
579            // nocov end
580        }))
581    }
582}
583
584/// Generate date strings in YYYY-MM-DD format.
585// nocov start
586pub fn dates() -> DateGenerator {
587    DateGenerator
588    // nocov end
589}
590
591/// Generator for time strings in HH:MM:SS format. Created by [`times()`].
592pub struct TimeGenerator;
593
594impl Generator<String> for TimeGenerator {
595    // nocov start
596    fn as_basic(&self) -> Option<BasicGenerator<'_, String>> {
597        Some(BasicGenerator::new(cbor_map! {"type" => "time"}, |raw| {
598            super::deserialize_value(raw)
599            // nocov end
600        }))
601    }
602}
603
604/// Generate time strings in HH:MM:SS format.
605// nocov start
606pub fn times() -> TimeGenerator {
607    TimeGenerator
608    // nocov end
609}
610
611/// Generator for ISO 8601 datetime strings. Created by [`datetimes()`].
612pub struct DateTimeGenerator;
613
614impl Generator<String> for DateTimeGenerator {
615    // nocov start
616    fn as_basic(&self) -> Option<BasicGenerator<'_, String>> {
617        Some(BasicGenerator::new(
618            cbor_map! {"type" => "datetime"},
619            super::deserialize_value,
620            // nocov end
621        ))
622    }
623}
624
625/// Generate ISO 8601 datetime strings.
626// nocov start
627pub fn datetimes() -> DateTimeGenerator {
628    DateTimeGenerator
629    // nocov end
630}
631
632/// Generator for UUID strings in canonical hyphenated form. Created by [`uuids()`].
633///
634/// By default generates UUIDs of any version. Use [`UuidsGenerator::version`]
635/// to restrict to a specific RFC 4122 version (1–5).
636pub struct UuidsGenerator {
637    version: Option<u8>,
638}
639
640impl UuidsGenerator {
641    /// Restrict to UUIDs of a specific version (1–5).
642    pub fn version(mut self, version: u8) -> Self {
643        self.version = Some(version);
644        self
645    }
646
647    fn build_schema(&self) -> Value {
648        match self.version {
649            Some(v) => cbor_map! {"type" => "uuid", "version" => v as u64},
650            None => cbor_map! {"type" => "uuid"},
651        }
652    }
653}
654
655impl Generator<String> for UuidsGenerator {
656    fn do_draw(&self, tc: &TestCase) -> String {
657        self.as_basic().unwrap().do_draw(tc)
658    }
659
660    fn as_basic(&self) -> Option<BasicGenerator<'_, String>> {
661        Some(BasicGenerator::new(
662            self.build_schema(),
663            super::deserialize_value,
664        ))
665    }
666}
667
668/// Generate UUID strings in canonical hyphenated form.
669///
670/// See [`UuidsGenerator`] for builder methods.
671pub fn uuids() -> UuidsGenerator {
672    UuidsGenerator { version: None }
673}
hegel/generators/strings.rs

hegel/generators/
strings.rs
