From 503f6e49c6d0f8cd775f1b73d46139cde232c248 Mon Sep 17 00:00:00 2001
From: Terry Cain <TerryCain@fico.com>
Date: Fri, 22 Sep 2017 13:36:20 +0100
Subject: [PATCH] Added option to pass in the factory to a decorated class

---
 README.rst       |  7 +++++++
 freezegun/api.py | 14 +++++++++-----
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/README.rst b/README.rst
index ea16e9bd..1e78d0f6 100644
--- a/README.rst
+++ b/README.rst
@@ -171,6 +171,13 @@ Freezegun allows moving time to specific dates.
             frozen_datetime.move_to(initial_datetime)
             assert frozen_datetime() == initial_datetime
 
+
+    @freeze_time("2012-01-14", as_arg=True)
+    def test(frozen_time):
+        assert datetime.datetime.now() == datetime.datetime(2012, 1, 14)
+        frozen_time.move_to("2014-02-12")
+        assert datetime.datetime.now() == datetime.datetime(2014, 2, 12)
+
 Parameter for ``move_to`` can be any valid ``freeze_time`` date (string, date, datetime).
 
 
diff --git a/freezegun/api.py b/freezegun/api.py
index e7776145..c6426354 100644
--- a/freezegun/api.py
+++ b/freezegun/api.py
@@ -305,7 +305,7 @@ def move_to(self, target_datetime):
 
 class _freeze_time(object):
 
-    def __init__(self, time_to_freeze_str, tz_offset, ignore, tick):
+    def __init__(self, time_to_freeze_str, tz_offset, ignore, tick, as_arg):
 
         self.time_to_freeze = _parse_time_to_freeze(time_to_freeze_str)
         self.tz_offset = tz_offset
@@ -313,6 +313,7 @@ def __init__(self, time_to_freeze_str, tz_offset, ignore, tick):
         self.tick = tick
         self.undo_changes = []
         self.modules_at_start = set()
+        self.as_arg = as_arg
 
     def __call__(self, func):
         if inspect.isclass(func):
@@ -496,8 +497,11 @@ def stop(self):
 
     def decorate_callable(self, func):
         def wrapper(*args, **kwargs):
-            with self:
-                result = func(*args, **kwargs)
+            with self as time_factory:
+                if self.as_arg:
+                    result = func(time_factory, *args, **kwargs)
+                else:
+                    result = func(*args, **kwargs)
             return result
         functools.update_wrapper(wrapper, func)
 
@@ -508,7 +512,7 @@ def wrapper(*args, **kwargs):
         return wrapper
 
 
-def freeze_time(time_to_freeze=None, tz_offset=0, ignore=None, tick=False):
+def freeze_time(time_to_freeze=None, tz_offset=0, ignore=None, tick=False, as_arg=False):
     # Python3 doesn't have basestring, but it does have str.
     try:
         string_type = basestring
@@ -534,7 +538,7 @@ def freeze_time(time_to_freeze=None, tz_offset=0, ignore=None, tick=False):
     ignore.append('django.utils.six.moves')
     ignore.append('threading')
     ignore.append('Queue')
-    return _freeze_time(time_to_freeze, tz_offset, ignore, tick)
+    return _freeze_time(time_to_freeze, tz_offset, ignore, tick, as_arg)
 
 
 # Setup adapters for sqlite