Bar feeds now support including extra columns besides OHLC values.

CHANGELOG

@@ -1,6 +1,7 @@
 Version 0.18 (TBD)
 . [NEW] Returns analyzer now store datetimes along with each sample. Thanks MatthiasKauer for implementing this.
 . [NEW] Quandl tool to build feed supports mapping column names.
+. [NEW] Bar feeds now support including extra columns besides OHLC values.
 . [FIX] Fixed zero division error in StochasticOscillator. Thanks Blake Jennings for reporting this.
 . [FIX] Fixed a rounding bug in the default fill strategy when calculating the volume left for a given instrument. Thanks Markus Trenkwalder for reporting this.
 . [FIX] Was not logging properly in backtesting broker when shares are not integers.

pyalgotrade/bar.py

@@ -113,6 +113,9 @@ def getPrice(self):
         """Returns the closing or adjusted closing price."""
         raise NotImplementedError()
 
+    def getExtraColumns(self):
+        return {}
+
 
 class BasicBar(Bar):
     # Optimization to reduce memory footprint.
@@ -125,10 +128,11 @@ class BasicBar(Bar):
         '__volume',
         '__adjClose',
         '__frequency',
-        '__useAdjustedValue'
+        '__useAdjustedValue',
+        '__extra',
     )
 
-    def __init__(self, dateTime, open_, high, low, close, volume, adjClose, frequency):
+    def __init__(self, dateTime, open_, high, low, close, volume, adjClose, frequency, extra={}):
         if high < low:
             raise Exception("high < low on %s" % (dateTime))
         elif high < open_:
@@ -149,6 +153,7 @@ def __init__(self, dateTime, open_, high, low, close, volume, adjClose, frequenc
         self.__adjClose = adjClose
         self.__frequency = frequency
         self.__useAdjustedValue = False
+        self.__extra = extra
 
     def __setstate__(self, state):
         (self.__dateTime,
@@ -159,7 +164,8 @@ def __setstate__(self, state):
             self.__volume,
             self.__adjClose,
             self.__frequency,
-            self.__useAdjustedValue) = state
+            self.__useAdjustedValue,
+            self.__extra) = state
 
     def __getstate__(self):
         return (
@@ -171,7 +177,8 @@ def __getstate__(self):
             self.__volume,
             self.__adjClose,
             self.__frequency,
-            self.__useAdjustedValue
+            self.__useAdjustedValue,
+            self.__extra
         )
 
     def setUseAdjustedValue(self, useAdjusted):
@@ -259,6 +266,8 @@ def getPrice(self):
         else:
             return self.__close
 
+    def getExtraColumns(self):
+        return self.__extra
 
 class Bars(object):
 

pyalgotrade/barfeed/__init__.py

@@ -126,7 +126,7 @@ def getLastBar(self, instrument):
         return self.__lastBars.get(instrument, None)
 
     def getDefaultInstrument(self):
-        """Returns the default instrument."""
+        """Returns the last instrument registered."""
         return self.__defaultInstrument
 
     def getRegisteredInstruments(self):

pyalgotrade/barfeed/csvfeed.py

@@ -139,6 +139,7 @@ def __init__(self, columnNames, dateTimeFormat, dailyBarTime, frequency, timezon
         self.__closeColName = columnNames["close"]
         self.__volumeColName = columnNames["volume"]
         self.__adjCloseColName = columnNames["adj_close"]
+        self.__columnNames = columnNames
 
     def _parseDate(self, dateString):
         ret = datetime.datetime.strptime(dateString, self.__dateTimeFormat)
@@ -173,7 +174,16 @@ def parseBar(self, csvRowDict):
             if len(adjCloseValue) > 0:
                 adjClose = float(adjCloseValue)
                 self.__haveAdjClose = True
-        return bar.BasicBar(dateTime, open_, high, low, close, volume, adjClose, self.__frequency)
+
+        # Extra columns
+        extra = {}
+        for k, v in csvRowDict.iteritems():
+            if k not in self.__columnNames:
+                extra[k] = csvutils.float_or_string(v)
+
+        return bar.BasicBar(
+            dateTime, open_, high, low, close, volume, adjClose, self.__frequency, extra=extra
+        )
 
 
 class GenericBarFeed(BarFeed):

pyalgotrade/dataseries/bards.py

@@ -38,8 +38,16 @@ def __init__(self, maxLen=None):
         self.__lowDS = dataseries.SequenceDataSeries(maxLen)
         self.__volumeDS = dataseries.SequenceDataSeries(maxLen)
         self.__adjCloseDS = dataseries.SequenceDataSeries(maxLen)
+        self.__extraDS = {}
         self.__useAdjustedValues = False
 
+    def __getOrCreateExtraDS(self, name):
+        ret = self.__extraDS.get(name)
+        if ret is None:
+            ret = dataseries.SequenceDataSeries(self.getMaxLen())
+            self.__extraDS[name] = ret
+        return ret
+
     def setUseAdjustedValues(self, useAdjusted):
         self.__useAdjustedValues = useAdjusted
 
@@ -50,14 +58,21 @@ def appendWithDateTime(self, dateTime, bar):
         assert(dateTime is not None)
         assert(bar is not None)
         bar.setUseAdjustedValue(self.__useAdjustedValues)
-        dataseries.SequenceDataSeries.appendWithDateTime(self, dateTime, bar)
+
+        super(BarDataSeries, self).appendWithDateTime(dateTime, bar)
+
         self.__openDS.appendWithDateTime(dateTime, bar.getOpen())
         self.__closeDS.appendWithDateTime(dateTime, bar.getClose())
         self.__highDS.appendWithDateTime(dateTime, bar.getHigh())
         self.__lowDS.appendWithDateTime(dateTime, bar.getLow())
         self.__volumeDS.appendWithDateTime(dateTime, bar.getVolume())
         self.__adjCloseDS.appendWithDateTime(dateTime, bar.getAdjClose())
 
+        # Process extra columns.
+        for name, value in bar.getExtraColumns().iteritems():
+            extraDS = self.__getOrCreateExtraDS(name)
+            extraDS.appendWithDateTime(dateTime, value)
+
     def getOpenDataSeries(self):
         """Returns a :class:`pyalgotrade.dataseries.DataSeries` with the open prices."""
         return self.__openDS
@@ -88,3 +103,8 @@ def getPriceDataSeries(self):
             return self.__adjCloseDS
         else:
             return self.__closeDS
+
+    def getExtraDataSeries(self, name):
+        """Returns a :class:`pyalgotrade.dataseries.DataSeries` for an extra column."""
+        return self.__getOrCreateExtraDS(name)
+

pyalgotrade/feed/csvfeed.py

@@ -129,11 +129,7 @@ def setTimeDelta(self, timeDelta):
 
 
 def float_or_string(column, value):
-    try:
-        ret = float(value)
-    except Exception:
-        ret = value
-    return ret
+    return csvutils.float_or_string(value)
 
 
 class Feed(BaseFeed):

pyalgotrade/utils/csvutils.py

@@ -68,3 +68,12 @@ def download_csv(url, url_params=None, content_type="text/csv"):
         ret = ret[1:]
 
     return ret
+
+
+def float_or_string(value):
+    try:
+        ret = float(value)
+    except Exception:
+        ret = value
+    return ret
+

testcases/quandl_test.py

@@ -180,3 +180,16 @@ def testMapColumnNames(self):
             self.assertEquals(bf["AAPL"][-1].getClose(), 322.560013)
             self.assertEquals(bf["AAPL"][-1].getAdjClose(), 42.674196)
             self.assertEquals(bf["AAPL"][-1].getPrice(), 42.674196 )
+
+    def testExtraColumns(self):
+        with common.TmpDir() as tmpPath:
+            columnNames = {
+                "open": "Last",
+                "close": "Last"
+            }
+            bf = quandl.build_feed("BITSTAMP", ["USD"], 2014, 2014, tmpPath, columnNames=columnNames)
+            bf.loadAll()
+            self.assertEquals(bf["USD"][-1].getExtraColumns()["Bid"], 319.19)
+            self.assertEquals(bf["USD"][-1].getExtraColumns()["Ask"], 319.63)
+            bids = bf["USD"].getExtraDataSeries("Bid")
+            self.assertEquals(bids[-1], 319.19)