summaryrefslogtreecommitdiff
path: root/factory
diff options
context:
space:
mode:
Diffstat (limited to 'factory')
-rw-r--r--factory/compat.py28
-rw-r--r--factory/fuzzy.py102
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)