Merge pull request #353 from WisdomPill/as_kwarg

Added as_kwarg argument

README.rst

@@ -43,6 +43,23 @@ Decorator
         def test_the_class(self):
             assert datetime.datetime.now() == datetime.datetime(2012, 1, 14)
 
+    # Or method decorator, might also pass frozen time object as kwarg
+
+    class TestUnitTestMethodDecorator(unittest.TestCase):
+        @freeze_time('2013-04-09')
+        def test_method_decorator_works_on_unittest(self):
+            self.assertEqual(datetime.date(2013, 4, 9), datetime.date.today())
+
+        @freeze_time('2013-04-09', as_kwarg='frozen_time')
+        def test_method_decorator_works_on_unittest(self, frozen_time):
+            self.assertEqual(datetime.date(2013, 4, 9), datetime.date.today())
+            self.assertEqual(datetime.date(2013, 4, 9), frozen_time.time_to_freeze.today())
+
+        @freeze_time('2013-04-09', as_kwarg='hello')
+        def test_method_decorator_works_on_unittest(self, **kwargs):
+            self.assertEqual(datetime.date(2013, 4, 9), datetime.date.today())
+            self.assertEqual(datetime.date(2013, 4, 9), kwargs.get('hello').time_to_freeze.today())
+
 Context manager
 ~~~~~~~~~~~~~~~
 

freezegun/api.py

@@ -513,8 +513,7 @@ def move_to(self, target_datetime):
 
 class _freeze_time(object):
 
-
-    def __init__(self, time_to_freeze_str, tz_offset, ignore, tick, as_arg, auto_tick_seconds):
+    def __init__(self, time_to_freeze_str, tz_offset, ignore, tick, as_arg, as_kwarg, auto_tick_seconds):
         self.time_to_freeze = _parse_time_to_freeze(time_to_freeze_str)
         self.tz_offset = _parse_tz_offset(tz_offset)
         self.ignore = tuple(ignore)
@@ -523,6 +522,7 @@ def __init__(self, time_to_freeze_str, tz_offset, ignore, tick, as_arg, auto_tic
         self.undo_changes = []
         self.modules_at_start = set()
         self.as_arg = as_arg
+        self.as_kwarg = as_kwarg
 
     def __call__(self, func):
         if inspect.isclass(func):
@@ -693,7 +693,7 @@ def stop(self):
                         continue
                     elif mod_name.startswith(self.ignore) or mod_name.endswith('.six.moves'):
                         continue
-                    elif (not hasattr(module, "__name__") or module.__name__ in ('datetime', 'time')):
+                    elif not hasattr(module, "__name__") or module.__name__ in ('datetime', 'time'):
                         continue
                     for module_attribute in dir(module):
 
@@ -729,8 +729,13 @@ def decorate_coroutine(self, coroutine):
     def decorate_callable(self, func):
         def wrapper(*args, **kwargs):
             with self as time_factory:
-                if self.as_arg:
+                if self.as_arg and self.as_kwarg:
+                    assert False, "You can't specify both as_arg and as_kwarg at the same time. Pick one."
+                elif self.as_arg:
                     result = func(time_factory, *args, **kwargs)
+                elif self.as_kwarg:
+                    kwargs[self.as_kwarg] = time_factory
+                    result = func(*args, **kwargs)
                 else:
                     result = func(*args, **kwargs)
             return result
@@ -743,7 +748,8 @@ def wrapper(*args, **kwargs):
         return wrapper
 
 
-def freeze_time(time_to_freeze=None, tz_offset=0, ignore=None, tick=False, as_arg=False, auto_tick_seconds=0):
+def freeze_time(time_to_freeze=None, tz_offset=0, ignore=None, tick=False, as_arg=False, as_kwarg='',
+                auto_tick_seconds=0):
     acceptable_times = (type(None), _string_type, datetime.date, datetime.timedelta,
              types.FunctionType, types.GeneratorType)
 
@@ -782,7 +788,7 @@ def freeze_time(time_to_freeze=None, tz_offset=0, ignore=None, tick=False, as_ar
         'gi',
     ])
 
-    return _freeze_time(time_to_freeze, tz_offset, ignore, tick, as_arg, auto_tick_seconds)
+    return _freeze_time(time_to_freeze, tz_offset, ignore, tick, as_arg, as_kwarg, auto_tick_seconds)
 
 
 # Setup adapters for sqlite

tests/test_datetimes.py

@@ -485,6 +485,22 @@ def test_isinstance_without_active():
     assert isinstance(today, datetime.date)
 
 
+class TestUnitTestMethodDecorator(unittest.TestCase):
+    @freeze_time('2013-04-09')
+    def test_method_decorator_works_on_unittest(self):
+        self.assertEqual(datetime.date(2013, 4, 9), datetime.date.today())
+
+    @freeze_time('2013-04-09', as_kwarg='frozen_time')
+    def test_method_decorator_works_on_unittest(self, frozen_time):
+        self.assertEqual(datetime.date(2013, 4, 9), datetime.date.today())
+        self.assertEqual(datetime.date(2013, 4, 9), frozen_time.time_to_freeze.today())
+
+    @freeze_time('2013-04-09', as_kwarg='hello')
+    def test_method_decorator_works_on_unittest(self, **kwargs):
+        self.assertEqual(datetime.date(2013, 4, 9), datetime.date.today())
+        self.assertEqual(datetime.date(2013, 4, 9), kwargs.get('hello').time_to_freeze.today())
+
+
 @freeze_time('2013-04-09')
 class TestUnitTestClassDecorator(unittest.TestCase):