Permalink
Browse files

fixing math operations per #2

  • Loading branch information...
1 parent d5b2c11 commit 8b101e398f2e980754551e20fb3f282d9b66bb3f @spulec committed Dec 14, 2012
Showing with 56 additions and 1 deletion.
  1. +25 −1 freezegun/api.py
  2. +31 −0 tests/test_operations.py
View
@@ -14,6 +14,18 @@ class FakeDate(real_date):
def __init__(self, *args, **kwargs):
return super(FakeDate, self).__init__(*args, **kwargs)
+ def __add__(self, other):
+ result = super(FakeDate, self).__add__(other)
+ if result is NotImplemented:
+ return result
+ return date_to_fakedate(result)
+
+ def __sub__(self, other):
+ result = super(FakeDate, self).__sub__(other)
+ if result is NotImplemented:
+ return result
+ return date_to_fakedate(result)
+
@classmethod
def today(cls):
if cls.active:
@@ -31,6 +43,18 @@ class FakeDatetime(real_datetime, FakeDate):
def __init__(self, *args, **kwargs):
return super(FakeDatetime, self).__init__(*args, **kwargs)
+ def __add__(self, other):
+ result = super(FakeDatetime, self).__add__(other)
+ if result is NotImplemented:
+ return result
+ return datetime_to_fakedatetime(result)
+
+ def __sub__(self, other):
+ result = super(FakeDatetime, self).__sub__(other)
+ if result is NotImplemented:
+ return result
+ return datetime_to_fakedatetime(result)
+
@classmethod
def now(cls):
if cls.active:
@@ -45,7 +69,7 @@ def utcnow(cls):
result = cls.time_to_freeze
else:
result = real_datetime.utcnow()
- return datetime_to_fakedatetime(result)
+ return result
datetime.datetime = FakeDatetime
datetime.date = FakeDate
@@ -0,0 +1,31 @@
+import datetime
+from freezegun import freeze_time
+from dateutil.relativedelta import relativedelta
+
+
+def test_addition():
+ now = datetime.datetime.now()
+ later = now + datetime.timedelta(days=1)
+ other_later = now + relativedelta(days=1)
+ assert isinstance(later, datetime.datetime)
+ assert isinstance(other_later, datetime.datetime)
+
+ today = datetime.date.today()
+ tomorrow = today + datetime.timedelta(days=1)
+ other_tomorrow = today + relativedelta(days=1)
+ assert isinstance(tomorrow, datetime.date)
+ assert isinstance(other_tomorrow, datetime.date)
+
+
+def test_subtraction():
+ now = datetime.datetime.now()
+ before = now - datetime.timedelta(days=1)
+ other_before = now - relativedelta(days=1)
+ assert isinstance(before, datetime.datetime)
+ assert isinstance(other_before, datetime.datetime)
+
+ today = datetime.date.today()
+ yesterday = today - datetime.timedelta(days=1)
+ other_yesterday = today - relativedelta(days=1)
+ assert isinstance(yesterday, datetime.date)
+ assert isinstance(other_yesterday, datetime.date)

0 comments on commit 8b101e3

Please sign in to comment.