File: constrained_strings.py

package info (click to toggle)
python-polyfactory 2.22.2-1
  • links: PTS, VCS
  • area: main
  • in suites: sid
  • size: 1,892 kB
  • sloc: python: 11,338; makefile: 103; sh: 37
file content (138 lines) | stat: -rw-r--r-- 3,900 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
from __future__ import annotations

from typing import TYPE_CHECKING, Callable, Pattern, TypeVar, Union, cast

from polyfactory.exceptions import ParameterException
from polyfactory.value_generators.primitives import create_random_bytes, create_random_string
from polyfactory.value_generators.regex import RegexFactory

T = TypeVar("T", bound=Union[bytes, str])

if TYPE_CHECKING:
    from random import Random


def _validate_length(
    min_length: int | None = None,
    max_length: int | None = None,
) -> None:
    """Validate the length parameters make sense.

    :param min_length: Minimum length.
    :param max_length: Maximum length.

    :raises: ParameterException.

    :returns: None.
    """
    if min_length is not None and min_length < 0:
        msg = "min_length must be greater or equal to 0"
        raise ParameterException(msg)

    if max_length is not None and max_length < 0:
        msg = "max_length must be greater or equal to 0"
        raise ParameterException(msg)

    if max_length is not None and min_length is not None and max_length < min_length:
        msg = "max_length must be greater than min_length"
        raise ParameterException(msg)


def _generate_pattern(
    random: Random,
    pattern: str | Pattern,
    lower_case: bool = False,
    upper_case: bool = False,
    min_length: int | None = None,
    max_length: int | None = None,
) -> str:
    """Generate a regex.

    :param random: An instance of random.
    :param pattern: A regex or string pattern.
    :param lower_case: Whether to lowercase the result.
    :param upper_case: Whether to uppercase the result.
    :param min_length: A minimum length.
    :param max_length: A maximum length.

    :returns: A string matching the given pattern.
    """
    regex_factory = RegexFactory(random=random)
    result = regex_factory(pattern)
    if min_length:
        while len(result) < min_length:
            result += regex_factory(pattern)

    if max_length is not None and len(result) > max_length:
        result = result[:max_length]

    if lower_case:
        result = result.lower()

    if upper_case:
        result = result.upper()

    return result


def handle_constrained_string_or_bytes(
    random: Random,
    t_type: Callable[[], T],
    lower_case: bool = False,
    upper_case: bool = False,
    min_length: int | None = None,
    max_length: int | None = None,
    pattern: str | Pattern | None = None,
) -> T:
    """Handle constrained string or bytes, for example - pydantic `constr` or `conbytes`.

    :param random: An instance of random.
    :param t_type: A type (str or bytes)
    :param lower_case: Whether to lowercase the result.
    :param upper_case: Whether to uppercase the result.
    :param min_length: A minimum length.
    :param max_length: A maximum length.
    :param pattern: A regex or string pattern.

    :returns: A value of type T.
    """
    _validate_length(min_length=min_length, max_length=max_length)

    if max_length == 0:
        return t_type()

    if pattern:
        return cast(
            "T",
            _generate_pattern(
                random=random,
                pattern=pattern,
                lower_case=lower_case,
                upper_case=upper_case,
                min_length=min_length,
                max_length=max_length,
            ),
        )

    if t_type is str:
        return cast(
            "T",
            create_random_string(
                min_length=min_length,
                max_length=max_length,
                lower_case=lower_case,
                upper_case=upper_case,
                random=random,
            ),
        )

    return cast(
        "T",
        create_random_bytes(
            min_length=min_length,
            max_length=max_length,
            lower_case=lower_case,
            upper_case=upper_case,
            random=random,
        ),
    )