ENH: Add currency column to EquityPricing.

Author: Scott Sanderson

tests/utils/test_numpy_utils.py

@@ -19,8 +19,10 @@
 from toolz import curry
 from toolz.curried.operator import ne
 
+from zipline.testing.predicates import assert_equal
 from zipline.utils.functional import mapall as lazy_mapall
 from zipline.utils.numpy_utils import (
+    bytes_array_to_native_str_object_array,
     is_float,
     is_int,
     is_datetime,
@@ -92,3 +94,13 @@ def test_is_datetime(self):
 
         for bad_value in everything_but(datetime, CASES):
             self.assertFalse(is_datetime(bad_value))
+
+
+class ArrayUtilsTestCase(TestCase):
+
+    def test_bytes_array_to_native_str_object_array(self):
+        a = array([b'abc', b'def'], dtype='S3')
+        result = bytes_array_to_native_str_object_array(a)
+        expected = array(['abc', 'def'], dtype=object)
+
+        assert_equal(result, expected)

zipline/pipeline/data/equity_pricing.py

@@ -1,7 +1,7 @@
 """
 Dataset representing OHLCV data.
 """
-from zipline.utils.numpy_utils import float64_dtype
+from zipline.utils.numpy_utils import float64_dtype, categorical_dtype
 
 from ..domain import US_EQUITIES
 from .dataset import Column, DataSet
@@ -17,6 +17,7 @@ class EquityPricing(DataSet):
     low = Column(float64_dtype)
     close = Column(float64_dtype)
     volume = Column(float64_dtype)
+    currency = Column(categorical_dtype)
 
 
 # Backwards compat alias.

zipline/pipeline/loaders/equity_pricing_loader.py

@@ -16,11 +16,17 @@
 from interface import implements
 from numpy import iinfo, uint32, multiply
 
+from zipline.currency import MISSING_CURRENCY_CODE
 from zipline.data.fx import ExplodingFXRateReader
 from zipline.lib.adjusted_array import AdjustedArray
+from zipline.utils.numpy_utils import (
+    repeat_first_axis,
+    bytes_array_to_native_str_object_array,
+)
 
 from .base import PipelineLoader
 from .utils import shift_dates
+from ..data.equity_pricing import EquityPricing
 
 UINT32_MAX = iinfo(uint32).max
 
@@ -81,9 +87,12 @@ def load_adjusted_array(self, domain, columns, dates, sids, mask):
         sessions = domain.all_sessions()
         shifted_dates = shift_dates(sessions, dates[0], dates[-1], shift=1)
 
-        colnames = [c.name for c in columns]
-        raw_arrays = self.raw_price_reader.load_raw_arrays(
-            colnames,
+        ohlcv_cols, currency_cols = self._split_column_types(columns)
+        del columns  # From here on we should use ohlcv_cols or currency_cols.
+        ohlcv_colnames = [c.name for c in ohlcv_cols]
+
+        raw_ohlcv_arrays = self.raw_price_reader.load_raw_arrays(
+            ohlcv_colnames,
             shifted_dates[0],
             shifted_dates[-1],
             sids,
@@ -93,25 +102,40 @@ def load_adjusted_array(self, domain, columns, dates, sids, mask):
         # dates to load currency conversion rates to make them line up with
         # dates used to fetch prices.
         self._inplace_currency_convert(
-            columns,
-            raw_arrays,
+            ohlcv_cols,
+            raw_ohlcv_arrays,
             shifted_dates,
             sids,
         )
 
         adjustments = self.adjustments_reader.load_pricing_adjustments(
-            colnames,
+            ohlcv_colnames,
             dates,
             sids,
         )
 
         out = {}
-        for c, c_raw, c_adjs in zip(columns, raw_arrays, adjustments):
+        for c, c_raw, c_adjs in zip(ohlcv_cols, raw_ohlcv_arrays, adjustments):
             out[c] = AdjustedArray(
                 c_raw.astype(c.dtype),
                 c_adjs,
                 c.missing_value,
             )
+
+        for c in currency_cols:
+            codes_1d = bytes_array_to_native_str_object_array(
+                self.raw_price_reader.currency_codes(sids)
+            )
+            # XXX: Should this just be the contract of `currency_codes`?
+            codes_1d[codes_1d == MISSING_CURRENCY_CODE] = None
+
+            codes = repeat_first_axis(codes_1d, len(dates))
+            out[c] = AdjustedArray(
+                codes,
+                adjustments={},
+                missing_value=None,
+            )
+
         return out
 
     @property
@@ -169,6 +193,33 @@ def _inplace_currency_convert(self, columns, arrays, dates, sids):
             for arr in arrays:
                 multiply(arr, rates, out=arr)
 
+    def _split_column_types(self, columns):
+        """Split out currency columns from OHLCV columns.
+
+        Parameters
+        ----------
+        columns : list[zipline.pipeline.data.BoundColumn]
+            Columns to be loaded by ``load_adjusted_array``.
+
+        Returns
+        -------
+        ohlcv_columns : list[zipline.pipeline.data.BoundColumn]
+            Price and volume columns from ``columns``.
+        currency_columns : list[zipline.pipeline.data.BoundColumn]
+            Currency code column from ``columns``, if present.
+        """
+        currency_name = EquityPricing.currency.name
+
+        ohlcv = []
+        currency = []
+        for c in columns:
+            if c.name == currency_name:
+                currency.append(c)
+            else:
+                ohlcv.append(c)
+
+        return ohlcv, currency
+
 
 # Backwards compat alias.
 USEquityPricingLoader = EquityPricingLoader

zipline/utils/numpy_utils.py

@@ -9,6 +9,7 @@
     filterwarnings,
 )
 
+import six
 import numpy as np
 from numpy import (
     array_equal,
@@ -502,3 +503,12 @@ def compare_datetime_arrays(x, y):
     """
 
     return array_equal(x.view('int64'), y.view('int64'))
+
+
+def bytes_array_to_native_str_object_array(a):
+    """Convert an array of dtype S to an object array containing `str`.
+    """
+    if six.PY2:
+        return a.astype(object)
+    else:
+        return a.astype(str).astype(object)