1# Copyright: See the LICENSE file.
2
3
4"""Additional declarations for "fuzzy" attribute definitions."""
5
6
7import datetime
8import decimal
9import string
10import warnings
11
12from . import declarations, random
13
14random_seed_warning = (
15    "Setting a specific random seed for {} can still have varying results "
16    "unless you also set a specific end date. For details and potential solutions "
17    "see https://github.com/FactoryBoy/factory_boy/issues/331"
18)
19
20
21class BaseFuzzyAttribute(declarations.BaseDeclaration):
22    """Base class for fuzzy attributes.
23
24    Custom fuzzers should override the `fuzz()` method.
25    """
26
27    def fuzz(self):  # pragma: no cover
28        raise NotImplementedError()
29
30    def evaluate(self, instance, step, extra):
31        return self.fuzz()
32
33
34class FuzzyAttribute(BaseFuzzyAttribute):
35    """Similar to LazyAttribute, but yields random values.
36
37    Attributes:
38        function (callable): function taking no parameters and returning a
39            random value.
40    """
41
42    def __init__(self, fuzzer):
43        super().__init__()
44        self.fuzzer = fuzzer
45
46    def fuzz(self):
47        return self.fuzzer()
48
49
50class FuzzyText(BaseFuzzyAttribute):
51    """Random string with a given prefix.
52
53    Generates a random string of the given length from chosen chars.
54    If a prefix or a suffix are supplied, they will be prepended / appended
55    to the generated string.
56
57    Args:
58        prefix (text): An optional prefix to prepend to the random string
59        length (int): the length of the random part
60        suffix (text): An optional suffix to append to the random string
61        chars (str list): the chars to choose from
62
63    Useful for generating unique attributes where the exact value is
64    not important.
65    """
66
67    def __init__(self, prefix='', length=12, suffix='', chars=string.ascii_letters):
68        super().__init__()
69        self.prefix = prefix
70        self.suffix = suffix
71        self.length = length
72        self.chars = tuple(chars)  # Unroll iterators
73
74    def fuzz(self):
75        chars = [random.randgen.choice(self.chars) for _i in range(self.length)]
76        return self.prefix + ''.join(chars) + self.suffix
77
78
79class FuzzyChoice(BaseFuzzyAttribute):
80    """Handles fuzzy choice of an attribute.
81
82    Args:
83        choices (iterable): An iterable yielding options; will only be unrolled
84            on the first call.
85        getter (callable or None): a function to parse returned values
86    """
87
88    def __init__(self, choices, getter=None):
89        self.choices = None
90        self.choices_generator = choices
91        self.getter = getter
92        super().__init__()
93
94    def fuzz(self):
95        if self.choices is None:
96            self.choices = list(self.choices_generator)
97        value = random.randgen.choice(self.choices)
98        if self.getter is None:
99            return value
100        return self.getter(value)
101
102
103class FuzzyInteger(BaseFuzzyAttribute):
104    """Random integer within a given range."""
105
106    def __init__(self, low, high=None, step=1):
107        if high is None:
108            high = low
109            low = 0
110
111        self.low = low
112        self.high = high
113        self.step = step
114
115        super().__init__()
116
117    def fuzz(self):
118        return random.randgen.randrange(self.low, self.high + 1, self.step)
119
120
121class FuzzyDecimal(BaseFuzzyAttribute):
122    """Random decimal within a given range."""
123
124    def __init__(self, low, high=None, precision=2):
125        if high is None:
126            high = low
127            low = 0.0
128
129        self.low = low
130        self.high = high
131        self.precision = precision
132
133        super().__init__()
134
135    def fuzz(self):
136        base = decimal.Decimal(str(random.randgen.uniform(self.low, self.high)))
137        return base.quantize(decimal.Decimal(10) ** -self.precision)
138
139
140class FuzzyFloat(BaseFuzzyAttribute):
141    """Random float within a given range."""
142
143    def __init__(self, low, high=None, precision=15):
144        if high is None:
145            high = low
146            low = 0
147
148        self.low = low
149        self.high = high
150        self.precision = precision
151
152        super().__init__()
153
154    def fuzz(self):
155        base = random.randgen.uniform(self.low, self.high)
156        return float(format(base, '.%dg' % self.precision))
157
158
159class FuzzyDate(BaseFuzzyAttribute):
160    """Random date within a given date range."""
161
162    def __init__(self, start_date, end_date=None):
163        super().__init__()
164        if end_date is None:
165            if random.randgen.state_set:
166                cls_name = self.__class__.__name__
167                warnings.warn(random_seed_warning.format(cls_name), stacklevel=2)
168            end_date = datetime.date.today()
169
170        if start_date > end_date:
171            raise ValueError(
172                "FuzzyDate boundaries should have start <= end; got %r > %r."
173                % (start_date, end_date))
174
175        self.start_date = start_date.toordinal()
176        self.end_date = end_date.toordinal()
177
178    def fuzz(self):
179        return datetime.date.fromordinal(random.randgen.randint(self.start_date, self.end_date))
180
181
182class BaseFuzzyDateTime(BaseFuzzyAttribute):
183    """Base class for fuzzy datetime-related attributes.
184
185    Provides fuzz() computation, forcing year/month/day/hour/...
186    """
187
188    def _check_bounds(self, start_dt, end_dt):
189        if start_dt > end_dt:
190            raise ValueError(
191                """%s boundaries should have start <= end, got %r > %r""" % (
192                    self.__class__.__name__, start_dt, end_dt))
193
194    def _now(self):
195        raise NotImplementedError()
196
197    def __init__(self, start_dt, end_dt=None,
198                 force_year=None, force_month=None, force_day=None,
199                 force_hour=None, force_minute=None, force_second=None,
200                 force_microsecond=None):
201        super().__init__()
202
203        if end_dt is None:
204            if random.randgen.state_set:
205                cls_name = self.__class__.__name__
206                warnings.warn(random_seed_warning.format(cls_name), stacklevel=2)
207            end_dt = self._now()
208
209        self._check_bounds(start_dt, end_dt)
210
211        self.start_dt = start_dt
212        self.end_dt = end_dt
213        self.force_year = force_year
214        self.force_month = force_month
215        self.force_day = force_day
216        self.force_hour = force_hour
217        self.force_minute = force_minute
218        self.force_second = force_second
219        self.force_microsecond = force_microsecond
220
221    def fuzz(self):
222        delta = self.end_dt - self.start_dt
223        microseconds = delta.microseconds + 1000000 * (delta.seconds + (delta.days * 86400))
224
225        offset = random.randgen.randint(0, microseconds)
226        result = self.start_dt + datetime.timedelta(microseconds=offset)
227
228        if self.force_year is not None:
229            result = result.replace(year=self.force_year)
230        if self.force_month is not None:
231            result = result.replace(month=self.force_month)
232        if self.force_day is not None:
233            result = result.replace(day=self.force_day)
234        if self.force_hour is not None:
235            result = result.replace(hour=self.force_hour)
236        if self.force_minute is not None:
237            result = result.replace(minute=self.force_minute)
238        if self.force_second is not None:
239            result = result.replace(second=self.force_second)
240        if self.force_microsecond is not None:
241            result = result.replace(microsecond=self.force_microsecond)
242
243        return result
244
245
246class FuzzyNaiveDateTime(BaseFuzzyDateTime):
247    """Random naive datetime within a given range.
248
249    If no upper bound is given, will default to datetime.datetime.now().
250    """
251
252    def _now(self):
253        return datetime.datetime.now()
254
255    def _check_bounds(self, start_dt, end_dt):
256        if start_dt.tzinfo is not None:
257            raise ValueError(
258                "FuzzyNaiveDateTime only handles naive datetimes, got start=%r"
259                % start_dt)
260        if end_dt.tzinfo is not None:
261            raise ValueError(
262                "FuzzyNaiveDateTime only handles naive datetimes, got end=%r"
263                % end_dt)
264        super()._check_bounds(start_dt, end_dt)
265
266
267class FuzzyDateTime(BaseFuzzyDateTime):
268    """Random timezone-aware datetime within a given range.
269
270    If no upper bound is given, will default to datetime.datetime.now()
271    If no timezone is given, will default to utc.
272    """
273
274    def _now(self):
275        return datetime.datetime.now(tz=datetime.timezone.utc)
276
277    def _check_bounds(self, start_dt, end_dt):
278        if start_dt.tzinfo is None:
279            raise ValueError(
280                "FuzzyDateTime requires timezone-aware datetimes, got start=%r"
281                % start_dt)
282        if end_dt.tzinfo is None:
283            raise ValueError(
284                "FuzzyDateTime requires timezone-aware datetimes, got end=%r"
285                % end_dt)
286        super()._check_bounds(start_dt, end_dt)
287