From 0b589ee39c372ab1e0a62b1381ef1a1b5b73f12d Mon Sep 17 00:00:00 2001
From: Thomas Wiecki <thomas.wiecki@gmail.com>
Date: Thu, 23 Aug 2012 16:03:02 -0400
Subject: [PATCH] Fixes to the updown algorithm. First unittest passes.

---
 tests/test_optimize.py         |  8 +++-----
 zipline/gens/tradegens.py      |  5 ++---
 zipline/lines.py               | 13 ++++++++++---
 zipline/optimize/algorithms.py | 10 +++++++---
 4 files changed, 22 insertions(+), 14 deletions(-)

diff --git a/tests/test_optimize.py b/tests/test_optimize.py
index 52e27bb8..83dca6a9 100644
--- a/tests/test_optimize.py
+++ b/tests/test_optimize.py
@@ -27,7 +27,7 @@ class TestUpDown(TestCase):
     def setUp(self):
         self.zipline_test_config = {
             'allocator'   : allocator,
-            'sid'         : 133,
+            'sid'         : [0],
             'trade_count' : 5,
             'amplitude'   : 30,
             'base_price'  : 50
@@ -48,7 +48,7 @@ class TestUpDown(TestCase):
 
         """
 
-        zipline, config = create_predictable_zipline(
+        algo, config = create_predictable_zipline(
             self.zipline_test_config,
             offset=0,
             simulate=False
@@ -74,9 +74,7 @@ class TestUpDown(TestCase):
             "Minimum price does not equal expected maximum price."
         )
 
-        zipline.run(config['trade_source'])
-
-        algo = config['algorithm']
+        algo.run(config['trade_source'])
 
         orders = np.asarray(algo.orders)
         max_order_idx = np.where(orders==orders.max())[0]
diff --git a/zipline/gens/tradegens.py b/zipline/gens/tradegens.py
index 5716f374..09685cba 100644
--- a/zipline/gens/tradegens.py
+++ b/zipline/gens/tradegens.py
@@ -182,7 +182,7 @@ class DataFrameSource(SpecificEquityTrades):
         self.data = data
         # Unpack config dictionary with default values.
         self.count = kwargs.get('count', 500)
-        self.sids = kwargs.get('sids', [1, 2])
+        self.sids = kwargs.get('sids', [0])
         self.start = kwargs.get('start', datetime(1957, 1, 1, 0, tzinfo = pytz.utc))
         self.end = kwargs.get('end', datetime(2010, 1, 1, tzinfo=pytz.utc))
         self.delta = kwargs.get('delta', timedelta(days = 1))
@@ -198,8 +198,7 @@ class DataFrameSource(SpecificEquityTrades):
     def create_fresh_generator(self):
         def _generator(df=self.data):
             for dt, series in df.iterrows():
-                dt = dt.tz_localize('UTC')
-                if (self.start > dt) or (dt < self.end):
+                if (dt < self.start) or (dt > self.end):
                     continue
                 event = {'dt': dt,
                          'source_id': self.get_hash(),
diff --git a/zipline/lines.py b/zipline/lines.py
index e4e4af8c..9f2a8153 100644
--- a/zipline/lines.py
+++ b/zipline/lines.py
@@ -179,7 +179,7 @@ class SimulatedTrading(object):
             else:
                 log.warning("Sending SIGINT")
                 os.kill(ppid, SIGINT)
-        
+
     def handle_exception(self, exc):
         if isinstance(exc, CancelSignal):
             # signal from monitor of an orderly shutdown,
@@ -206,7 +206,7 @@ class SimulatedTrading(object):
                     exc_type.__name__,
                     exc_value.message
                 )
-            
+
             self.results_socket.send(msg)
         except:
             log.exception("Exception while reporting simulation exception.")
@@ -380,6 +380,13 @@ class SimulatedTradingLite(object):
         self.with_tnfms = sequential_transforms(self.date_sorted, *self.transforms)
         self.trading_client = tsc(algorithm, environment, style)
         self.gen = self.trading_client.simulate(self.with_tnfms)
-        
+
+
     def get_results(self):
         return self.gen
+
+    def __iter__(self):
+        return self
+
+    def next(self):
+        return self.gen.next()
diff --git a/zipline/optimize/algorithms.py b/zipline/optimize/algorithms.py
index dc80ef5a..727e43db 100644
--- a/zipline/optimize/algorithms.py
+++ b/zipline/optimize/algorithms.py
@@ -10,8 +10,12 @@ from zipline.protocol import DATASOURCE_TYPE
 from zipline import ndict
 from zipline.utils.factory import create_trading_environment
 from zipline.gens.transform import StatefulTransform
-from zipline.lines import SimulatedTradingLite
+from zipline.lines import SimulatedTradingLite,  SimulatedTrading
 
+from logbook import Logger
+
+
+logger = Logger('Algo')
 class BuySellAlgorithm(object):
     """Algorithm that buys and sells alternatingly. The amount for
     each order can be specified. In addition, an offset that will
@@ -118,7 +122,7 @@ class TradingAlgorithm(object):
         self._setup(compute_risk_metrics=compute_risk_metrics)
 
         # drain simulated_trading
-        perfs = [perf for perf in self.simulated_trading]
+        perfs = list(self.simulated_trading)
 
         daily_stats = self._create_daily_stats(perfs)
         return daily_stats
@@ -174,7 +178,7 @@ class BuySellAlgorithmNew(TradingAlgorithm):
 
     def handle_data(self, data):
         order_size = self.buy_or_sell * (self.amount - (self.offset**2))
-        self.order(self.sid, order_size)
+        self.order(self.sids[0], order_size)
 
         #sell next time around.
         self.buy_or_sell *= -1