From daf29c2d39d9b6229f1345f464141f0ac93f8513 Mon Sep 17 00:00:00 2001
From: Thomas Wiecki <thomas.wiecki@gmail.com>
Date: Thu, 20 Dec 2012 12:49:24 -0500
Subject: [PATCH] ENH: Add granularity and annualizer arguments to
 TradingAlgorithm. Accompanying doc string and unittest.

---
 tests/test_algorithm.py | 14 ++++++++++++++
 zipline/algorithm.py    | 36 ++++++++++++++++++++++++++++++++----
 2 files changed, 46 insertions(+), 4 deletions(-)

diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py
index d757c04a..6485ccce 100644
--- a/tests/test_algorithm.py
+++ b/tests/test_algorithm.py
@@ -89,3 +89,17 @@ class TestTransformAlgorithm(TestCase):
         assert algo.registered_transforms['mavg']['kwargs'] == \
             {'window_length': 2, 'market_aware': True}
         assert algo.registered_transforms['mavg']['class'] is MovingAverage
+
+    def test_granularity_setting(self):
+        algo = TestRegisterTransformAlgorithm(granularity='daily')
+        self.assertEqual(algo.granularity, 'daily')
+        self.assertEqual(algo.annualizer, 250)
+
+        algo = TestRegisterTransformAlgorithm(granularity='minutely')
+        self.assertEqual(algo.granularity, 'minutely')
+        self.assertEqual(algo.annualizer, 250 * 6 * 60)
+
+        algo = TestRegisterTransformAlgorithm(granularity='minutely',
+                                              annualizer=10)
+        self.assertEqual(algo.granularity, 'minutely')
+        self.assertEqual(algo.annualizer, 10)
diff --git a/zipline/algorithm.py b/zipline/algorithm.py
index dda384fb..7622c665 100644
--- a/zipline/algorithm.py
+++ b/zipline/algorithm.py
@@ -65,8 +65,16 @@ class TradingAlgorithm(object):
 
     """
     def __init__(self, *args, **kwargs):
-        """
-        Initialize sids and other state variables.
+        """Initialize sids and other state variables.
+
+        :Arguments:
+            granularity : str (daily, hourly or minutely)
+               The duration of the bars.
+            annualizer : int <optional>
+               Which constant to use for annualizing risk metrics.
+               If not provided, will extract from granularity.
+            capital_base : float <default: 1.0e5>
+               How much capital to start with.
         """
         self.done = False
         self.order = None
@@ -84,16 +92,36 @@ class TradingAlgorithm(object):
         self.slippage = VolumeShareSlippage()
         self.commission = PerShare()
 
+        self.granularity = kwargs.get('granularity', 'daily')
+        # annualizer is used for e.g. risk calculations
+        self.annualizer = kwargs.get('annualizer', None)
+
         # set the capital base
         self.capital_base = kwargs.get('capital_base', DEFAULT_CAPITAL_BASE)
 
-        # an algorithm subclass needs to set initialized to True
-        # when it is fully initialized.
+        # an algorithm subclass needs to set initialized to True when
+        # it is fully initialized.
         self.initialized = False
 
         # call to user-defined constructor method
         self.initialize(*args, **kwargs)
 
+        # set annualizer according to granularity
+        # this is happening after initialize because granularity
+        # could be set in there.
+        if self.annualizer is None:
+            if self.granularity == 'daily':
+                self.annualizer = 250
+            elif self.granularity == 'hourly':
+                # trading days * hours
+                self.annualizer = 250 * 6
+            elif self.granularity == 'minutely':
+                # trading days * hours * minutes
+                self.annualizer = 250 * 6 * 60
+            else:
+                raise NotImplementedError('{g} is not implemented.\
+                                          '.format(g=self.granularity))
+
     def _create_generator(self, environment):
         """
         Create a basic generator setup using the sources and