From 4ba35d7d466d2d0f1934c2a97567c89fb7e5bf7c Mon Sep 17 00:00:00 2001
From: Richard Frank <rich@quantopian.com>
Date: Fri, 19 Apr 2013 12:22:06 -0400
Subject: [PATCH] ENH: Stream benchmark and treasury data when downloading

Instead of loading entire csv or xml into memory.
---
 zipline/data/benchmarks.py | 11 +++---
 zipline/data/treasuries.py | 71 ++++++++++++++++++++++++++++++--------
 2 files changed, 61 insertions(+), 21 deletions(-)

diff --git a/zipline/data/benchmarks.py b/zipline/data/benchmarks.py
index 1338bdc8..f1c645a6 100644
--- a/zipline/data/benchmarks.py
+++ b/zipline/data/benchmarks.py
@@ -18,7 +18,6 @@ from datetime import datetime
 
 import csv
 
-from StringIO import StringIO
 from functools import partial
 
 import requests
@@ -75,7 +74,7 @@ def get_raw_benchmark_data(start_date, end_date, symbol):
     ))
 
     res = requests.get('http://ichart.yahoo.com/table.csv',
-                       params=params)
+                       params=params, stream=True)
 
     if not res.ok:
         raise BenchmarkDataNotFoundError("""
@@ -85,7 +84,7 @@ start_date={start_date}, end_date={end_date}, url={url}""".strip().
                                                 end_date=end_date,
                                                 url=res.url))
 
-    return csv.DictReader(StringIO(res.content))
+    return csv.DictReader(res.iter_lines())
 
 
 def get_benchmark_data(symbol, start_date=None, end_date=None):
@@ -98,12 +97,10 @@ def get_benchmark_data(symbol, start_date=None, end_date=None):
         end_date = datetime.utcnow()
 
     raw_benchmark_data = get_raw_benchmark_data(start_date, end_date, symbol)
-    # Reverse data so we can load it in reverse chron order.
-    benchmarks_source = reversed(list(raw_benchmark_data))
 
     mappings = benchmark_mappings()
 
-    return source_to_records(mappings, benchmarks_source)
+    return source_to_records(mappings, raw_benchmark_data)
 
 
 def get_benchmark_returns(symbol, start_date=None, end_date=None):
@@ -120,4 +117,6 @@ def get_benchmark_returns(symbol, start_date=None, end_date=None):
         daily_return = DailyReturn(date=data_point['date'], returns=returns)
         benchmark_returns.append(daily_return)
 
+    # Reverse data so we can load it in reverse chron order.
+    benchmark_returns.reverse()
     return benchmark_returns
diff --git a/zipline/data/treasuries.py b/zipline/data/treasuries.py
index e46c6986..9a895ab8 100644
--- a/zipline/data/treasuries.py
+++ b/zipline/data/treasuries.py
@@ -16,8 +16,8 @@
 
 import requests
 
-from StringIO import StringIO
-from xml.dom.minidom import parse
+from collections import OrderedDict
+import xml.etree.ElementTree as ET
 
 from loader_utils import (
     guarded_conversion,
@@ -62,26 +62,67 @@ def treasury_mappings():
             in _CURVE_MAPPINGS.iteritems()}
 
 
+class iter_to_stream(object):
+    """
+    Exposes an iterable as an i/o stream
+    """
+    def __init__(self, iterable):
+        self.buffered = ""
+        self.iter = iter(iterable)
+
+    def read(self, size):
+        result = ""
+        while size > 0:
+            data = self.buffered or next(self.iter, None)
+            self.buffered = ""
+            if data is None:
+                break
+            size -= len(data)
+            if size < 0:
+                data, self.buffered = data[:size], data[size:]
+            result += data
+        return result
+
+
 def get_treasury_source():
     url = """\
 http://data.treasury.gov/feed.svc/DailyTreasuryYieldCurveRateData\
 """
-    res = requests.get(url)
+    res = requests.get(url, stream=True)
+    stream = iter_to_stream(res.iter_lines())
 
-    content = StringIO(res.content)
-    dom = parse(content)
+    elements = ET.iterparse(stream, ('end', 'start-ns', 'end-ns'))
 
-    entries = dom.getElementsByTagName("entry")
+    namespaces = OrderedDict()
+    properties_xpath = ['']
 
-    for entry in entries:
-        properties = entry.getElementsByTagName("m:properties")
-        datum = {node.nodeName.replace('d:', ''):
-                 node.childNodes[0].nodeValue
-                 if len(node.childNodes)
-                 else None
-                 for node in properties[0].childNodes
-                 if node.nodeType == dom.ELEMENT_NODE}
-        yield datum
+    def updated_namespaces():
+        if '' in namespaces and 'm' in namespaces:
+            properties_xpath[0] = "{%s}content/{%s}properties" % (
+                namespaces[''], namespaces['m']
+            )
+        else:
+            properties_xpath[0] = ''
+
+    for event, element in elements:
+        if event == 'end':
+            tag = ET.QName(element.tag).localname
+            if tag == "entry":
+                properties = element.find(properties_xpath[0])
+                datum = {ET.QName(node.tag).localname: node.text
+                         for node in properties.iterchildren()
+                         if ET.iselement(node)}
+                # clear the element after we've dealt with it:
+                element.clear()
+                yield datum
+
+        elif event == 'start-ns':
+            namespaces[element[0]] = element[1]
+            updated_namespaces()
+
+        elif event == 'end-ns':
+            namespaces.popitem()
+            updated_namespaces()
 
 
 def get_treasury_data():