diff options
Diffstat (limited to 'factory')
-rw-r--r-- | factory/compat.py | 28 | ||||
-rw-r--r-- | factory/fuzzy.py | 102 |
2 files changed, 130 insertions, 0 deletions
diff --git a/factory/compat.py b/factory/compat.py index 84f31b7..9594550 100644 --- a/factory/compat.py +++ b/factory/compat.py @@ -23,6 +23,7 @@ """Compatibility tools""" +import datetime import sys PY2 = (sys.version_info[0] == 2) @@ -33,3 +34,30 @@ if PY2: else: def is_string(obj): return isinstance(obj, str) + +try: + # Python >= 3.2 + UTC = datetime.timezone.utc +except AttributeError: + try: + # Fallback to pytz + from pytz import UTC + except ImportError: + + # Ok, let's write our own. + class _UTC(datetime.tzinfo): + """The UTC tzinfo.""" + + def utcoffset(self, dt): + return datetime.timedelta(0) + + def tzname(self, dt): + return "UTC" + + def dst(self, dt): + return datetime.timedelta(0) + + def localize(self, dt): + dt.astimezone(self) + + UTC = _UTC() diff --git a/factory/fuzzy.py b/factory/fuzzy.py index fea7b05..e266b93 100644 --- a/factory/fuzzy.py +++ b/factory/fuzzy.py @@ -27,6 +27,7 @@ import random import datetime +from . import compat from . import declarations @@ -104,3 +105,104 @@ class FuzzyDate(BaseFuzzyAttribute): def fuzz(self): return datetime.date.fromordinal(random.randint(self.start_date, self.end_date)) + + +class BaseFuzzyDateTime(BaseFuzzyAttribute): + """Base class for fuzzy datetime-related attributes. + + Provides fuzz() computation, forcing year/month/day/hour/... + """ + + def _check_bounds(self, start_dt, end_dt): + if start_dt > end_dt: + raise ValueError( + """%s boundaries should have start <= end, got %r > %r""" % ( + self.__class__.__name__, start_dt, end_dt)) + + def __init__(self, start_dt, end_dt=None, + force_year=None, force_month=None, force_day=None, + force_hour=None, force_minute=None, force_second=None, + force_microsecond=None, **kwargs): + super(BaseFuzzyDateTime, self).__init__(**kwargs) + + if end_dt is None: + end_dt = self._now() + + self._check_bounds(start_dt, end_dt) + + self.start_dt = start_dt + self.end_dt = end_dt + self.force_year = force_year + self.force_month = force_month + self.force_day = force_day + self.force_hour = force_hour + self.force_minute = force_minute + self.force_second = force_second + self.force_microsecond = force_microsecond + + def fuzz(self): + delta = self.end_dt - self.start_dt + microseconds = delta.microseconds + 1000000 * (delta.seconds + (delta.days * 86400)) + + offset = random.randint(0, microseconds) + result = self.start_dt + datetime.timedelta(microseconds=offset) + + if self.force_year is not None: + result = result.replace(year=self.force_year) + if self.force_month is not None: + result = result.replace(month=self.force_month) + if self.force_day is not None: + result = result.replace(day=self.force_day) + if self.force_hour is not None: + result = result.replace(hour=self.force_hour) + if self.force_minute is not None: + result = result.replace(minute=self.force_minute) + if self.force_second is not None: + result = result.replace(second=self.force_second) + if self.force_microsecond is not None: + result = result.replace(microsecond=self.force_microsecond) + + return result + + +class FuzzyNaiveDateTime(BaseFuzzyDateTime): + """Random naive datetime within a given range. + + If no upper bound is given, will default to datetime.datetime.utcnow(). + """ + + def _now(self): + return datetime.datetime.now() + + def _check_bounds(self, start_dt, end_dt): + if start_dt.tzinfo is not None: + raise ValueError( + "FuzzyNaiveDateTime only handles naive datetimes, got start=%r" + % start_dt) + if end_dt.tzinfo is not None: + raise ValueError( + "FuzzyNaiveDateTime only handles naive datetimes, got end=%r" + % end_dt) + super(FuzzyNaiveDateTime, self)._check_bounds(start_dt, end_dt) + + +class FuzzyDateTime(BaseFuzzyDateTime): + """Random timezone-aware datetime within a given range. + + If no upper bound is given, will default to datetime.datetime.now() + If no timezone is given, will default to utc. + """ + + def _now(self): + return datetime.datetime.now(tz=compat.UTC) + + def _check_bounds(self, start_dt, end_dt): + if start_dt.tzinfo is None: + raise ValueError( + "FuzzyDateTime only handles aware datetimes, got start=%r" + % start_dt) + if end_dt.tzinfo is None: + raise ValueError( + "FuzzyDateTime only handles aware datetimes, got end=%r" + % end_dt) + super(FuzzyDateTime, self)._check_bounds(start_dt, end_dt) |