Add FuzzyDateTime/FuzzyNaiveDateTime.

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)