From fa427afb16e58703a45e1bf1d106a4fc0d473872 Mon Sep 17 00:00:00 2001
From: fawce <fawce@quantopian.com>
Date: Mon, 29 Oct 2012 17:22:42 -0400
Subject: [PATCH] added algotime get/set to algorithm

---
 zipline/algorithm.py            | 21 +++++++++++++++++++++
 zipline/gens/tradesimulation.py |  1 +
 2 files changed, 22 insertions(+)

diff --git a/zipline/algorithm.py b/zipline/algorithm.py
index 33fca5ac..3d6b26b6 100644
--- a/zipline/algorithm.py
+++ b/zipline/algorithm.py
@@ -12,10 +12,14 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from copy import copy
 
+import pytz
 import pandas as pd
 import numpy as np
 
+from datetime import datetime
+
 from zipline.sources import DataFrameSource
 from zipline.utils.factory import create_trading_environment
 from zipline.transforms.utils import StatefulTransform
@@ -63,6 +67,7 @@ class TradingAlgorithm(object):
         self.order = None
         self.frame_count = 0
         self.portfolio = None
+        self.datetime = None
 
         self.registered_transforms = {}
         self.transforms = []
@@ -221,6 +226,22 @@ class TradingAlgorithm(object):
     def set_logger(self, logger):
         self.logger = logger
 
+    def set_datetime(self, dt):
+        assert isinstance(dt, datetime), \
+            "Attempt to set algorithm's current time with non-datetime"
+        assert dt.tzinfo == pytz.utc, \
+            "Algorithm expects a utc datetime"
+        self.datetime = dt
+
+    def get_datetime(self):
+        """
+        Returns a copy of the datetime.
+        """
+        date_copy = copy(self.datetime)
+        assert date_copy.tzinfo == pytz.utc, \
+            "Algorithm should have a utc datetime"
+        return date_copy
+
     def init(self, *args, **kwargs):
         """Called from constructor."""
         pass
diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py
index 79bc3f02..ce711487 100644
--- a/zipline/gens/tradesimulation.py
+++ b/zipline/gens/tradesimulation.py
@@ -269,6 +269,7 @@ class AlgorithmSimulator(object):
         # Needs to be set so that we inject the proper date into algo
         # log/print lines.
         self.snapshot_dt = date
+        self.algo.set_datetime(self.snapshot_dt)
         start_tic = datetime.now()
         self.algo.handle_data(self.universe)
         stop_tic = datetime.now()