diff --git a/zipline/utils/factory.py b/zipline/utils/factory.py index 57b3eaad..2e89ea21 100644 --- a/zipline/utils/factory.py +++ b/zipline/utils/factory.py @@ -377,18 +377,86 @@ def _load_raw_yahoo_data(indexes=None, stocks=None, start=None, end=None): return data -def load_from_yahoo(indexes=None, stocks=None, start=None, end=None): +def load_from_yahoo(indexes=None, + stocks=None, + start=None, + end=None, + adjusted=True): + """ + Loads price data from Yahoo into a dataframe for each of the indicated + securities. By default, 'price' is taken from Yahoo's 'Adjusted Close', + which removes the impact of splits and dividends. If the argument + 'adjusted' is False, then the non-adjusted 'close' field is used instead. + + :Arguments: + indexes : dict (Default: {'SPX': '^GSPC'}) + Financial indexes to load. + stocks : list (Default: ['AAPL', 'GE', 'IBM', 'MSFT', + 'XOM', 'AA', 'JNJ', 'PEP', 'KO']) + Stock closing prices to load. + start : datetime (Default: datetime(1993, 1, 1, 0, 0, 0, 0, pytz.utc)) + Retrieve prices from start date on. + end : datetime (Default: datetime(2002, 1, 1, 0, 0, 0, 0, pytz.utc)) + Retrieve prices until end date. + adjusted : bool (Default: True) + Adjust the price for splits and dividends. + + """ data = _load_raw_yahoo_data(indexes, stocks, start, end) - df = pd.DataFrame({key: d['Adj Close'] for key, d in data.iteritems()}) + if adjusted: + close_key = 'Adj Close' + else: + close_key = 'Close' + df = pd.DataFrame({key: d[close_key] for key, d in data.iteritems()}) df.index = df.index.tz_localize(pytz.utc) return df -def load_bars_from_yahoo(indexes=None, stocks=None, start=None, end=None): +def load_bars_from_yahoo(indexes=None, + stocks=None, + start=None, + end=None, + adjusted=True): + """ + Loads data from Yahoo into a panel with the following + column names for each indicated security: + - open + - high + - low + - close + - volume + - price + + Note that 'price' is Yahoo's 'Adjusted Close', which removes the + impact of splits and dividends. If the argument 'adjusted' is True, then + the open, high, low, and close values are adjusted as well. + + :Arguments: + indexes : dict (Default: {'SPX': '^GSPC'}) + Financial indexes to load. + stocks : list (Default: ['AAPL', 'GE', 'IBM', 'MSFT', + 'XOM', 'AA', 'JNJ', 'PEP', 'KO']) + Stock closing prices to load. + start : datetime (Default: datetime(1993, 1, 1, 0, 0, 0, 0, pytz.utc)) + Retrieve prices from start date on. + end : datetime (Default: datetime(2002, 1, 1, 0, 0, 0, 0, pytz.utc)) + Retrieve prices until end date. + adjusted : bool (Default: True) + Adjust open/high/low/close for splits and dividends. The 'price' + field is always adjusted. + + """ data = _load_raw_yahoo_data(indexes, stocks, start, end) panel = pd.Panel(data) # Rename columns panel.minor_axis = ['open', 'high', 'low', 'close', 'volume', 'price'] panel.major_axis = panel.major_axis.tz_localize(pytz.utc) - + # Adjust data + if adjusted: + adj_cols = ['open', 'high', 'low', 'close'] + for ticker in panel.items: + ratio = (panel[ticker]['price'] / panel[ticker]['close']) + ratio_filtered = ratio.fillna(0).values + for col in adj_cols: + panel[ticker][col] *= ratio_filtered return panel