from six import StringIO, iteritems from abc import ABCMeta, abstractmethod from collections import namedtuple import hashlib from textwrap import dedent import pandas as pd from pandas import read_csv import numpy from logbook import Logger import pytz import warnings import requests from zipline.errors import ( MultipleSymbolsFound, SymbolNotFound, ZiplineError ) from zipline.protocol import ( DATASOURCE_TYPE, Event ) from zipline.assets import Equity logger = Logger('Requests Source Logger') def roll_dts_to_midnight(dts, env): if len(dts) == 0: return dts return pd.DatetimeIndex( (dts.tz_convert('US/Eastern') - pd.Timedelta(hours=16)).date, tz='UTC', ) + env.trading_day class FetcherEvent(Event): pass class FetcherCSVRedirectError(ZiplineError): msg = dedent( """\ Attempt to fetch_csv from a redirected url. {url} must be changed to {new_url} """ ) def __init__(self, *args, **kwargs): self.url = kwargs["url"] self.new_url = kwargs["new_url"] self.extra = kwargs["extra"] super(FetcherCSVRedirectError, self).__init__(*args, **kwargs) # The following optional arguments are supported for # requests backed data sources. # see http://docs.python-requests.org/en/latest/api/#main-interface # for a full list. ALLOWED_REQUESTS_KWARGS = { 'params', 'headers', 'auth', 'cert'} # The following optional arguments are supported for pandas' read_csv # function, and may be passed as kwargs to the datasource below. # see http://pandas.pydata.org/ # pandas-docs/stable/generated/pandas.io.parsers.read_csv.html ALLOWED_READ_CSV_KWARGS = { 'sep', 'dialect', 'doublequote', 'escapechar', 'quotechar', 'quoting', 'skipinitialspace', 'lineterminator', 'header', 'index_col', 'names', 'prefix', 'skiprows', 'skipfooter', 'skip_footer', 'na_values', 'true_values', 'false_values', 'delimiter', 'converters', 'dtype', 'delim_whitespace', 'as_recarray', 'na_filter', 'compact_ints', 'use_unsigned', 'buffer_lines', 'warn_bad_lines', 'error_bad_lines', 'keep_default_na', 'thousands', 'comment', 'decimal', 'keep_date_col', 'nrows', 'chunksize', 'encoding', 'usecols' } SHARED_REQUESTS_KWARGS = { 'stream': True, 'allow_redirects': False, } def mask_requests_args(url, validating=False, params_checker=None, **kwargs): requests_kwargs = {key: val for (key, val) in iteritems(kwargs) if key in ALLOWED_REQUESTS_KWARGS} if params_checker is not None: url, s_params = params_checker(url) if s_params: if 'params' in requests_kwargs: requests_kwargs['params'].update(s_params) else: requests_kwargs['params'] = s_params # Giving the connection 30 seconds. This timeout does not # apply to the download of the response body. # (Note that Quandl links can take >10 seconds to return their # first byte on occasion) requests_kwargs['timeout'] = 1.0 if validating else 30.0 requests_kwargs.update(SHARED_REQUESTS_KWARGS) request_pair = namedtuple("RequestPair", ("requests_kwargs", "url")) return request_pair(requests_kwargs, url) class PandasCSV(object): __metaclass__ = ABCMeta def __init__(self, pre_func, post_func, env, start_date, end_date, date_column, date_format, timezone, symbol, mask, symbol_column, data_frequency, **kwargs): self.start_date = start_date self.end_date = end_date self.date_column = date_column self.date_format = date_format self.timezone = timezone self.mask = mask self.symbol_column = symbol_column or "symbol" self.data_frequency = data_frequency invalid_kwargs = set(kwargs) - ALLOWED_READ_CSV_KWARGS if invalid_kwargs: raise TypeError( "Unexpected keyword arguments: %s" % invalid_kwargs, ) self.pandas_kwargs = self.mask_pandas_args(kwargs) self.symbol = symbol self.env = env self.finder = env.asset_finder self.pre_func = pre_func self.post_func = post_func @property def fields(self): return self.df.columns.tolist() def get_hash(self): return self.namestring @abstractmethod def fetch_data(self): return @staticmethod def parse_date_str_series(format_str, tz, date_str_series, data_frequency, env): """ Efficient parsing for a 1d Pandas/numpy object containing string representations of dates. Note: pd.to_datetime is significantly faster when no format string is passed, and in pandas 0.12.0 the %p strptime directive is not correctly handled if a format string is explicitly passed, but AM/PM is handled properly if format=None. Moreover, we were previously ignoring this parameter unintentionally because we were incorrectly passing it as a positional. For all these reasons, we ignore the format_str parameter when parsing datetimes. """ # Explicitly ignoring this parameter. See note above. if format_str is not None: logger.warn( "The 'format_str' parameter to fetch_csv is deprecated. " "Ignoring and defaulting to pandas default date parsing." ) format_str = None tz_str = str(tz) if tz_str == pytz.utc.zone: parsed = pd.to_datetime( date_str_series.values, format=format_str, utc=True, coerce=True, ) else: parsed = pd.to_datetime( date_str_series.values, format=format_str, coerce=True, ).tz_localize(tz_str).tz_convert('UTC') if data_frequency == 'daily': parsed = roll_dts_to_midnight(parsed, env) return parsed def mask_pandas_args(self, kwargs): pandas_kwargs = {key: val for (key, val) in iteritems(kwargs) if key in ALLOWED_READ_CSV_KWARGS} if 'usecols' in pandas_kwargs: usecols = pandas_kwargs['usecols'] if usecols and self.date_column not in usecols: # make a new list so we don't modify user's, # and to ensure it is mutable with_date = list(usecols) with_date.append(self.date_column) pandas_kwargs['usecols'] = with_date # No strings in the 'symbol' column should be interpreted as NaNs pandas_kwargs.setdefault('keep_default_na', False) pandas_kwargs.setdefault('na_values', {'symbol': []}) return pandas_kwargs def _lookup_unconflicted_symbol(self, symbol): """ Attempt to find a unique asset whose symbol is the given string. If multiple assets have held the given symbol, return a 0. If no asset has held the given symbol, return a NaN. """ try: uppered = symbol.upper() except AttributeError: # The mapping fails because symbol was a non-string return numpy.nan try: return self.finder.lookup_symbol(uppered, as_of_date=None) except MultipleSymbolsFound: # Fill conflicted entries with zeros to mark that they need to be # resolved by date. return 0 except SymbolNotFound: # Fill not found entries with nans. return numpy.nan def load_df(self): df = self.fetch_data() if self.pre_func: df = self.pre_func(df) # Batch-convert the user-specifed date column into timestamps. df['dt'] = self.parse_date_str_series( self.date_format, self.timezone, df[self.date_column], self.data_frequency, self.env ).values # ignore rows whose dates we couldn't parse df = df[df['dt'].notnull()] if self.symbol is not None: df['sid'] = self.symbol elif self.finder: df.sort(self.symbol_column) # Pop the 'sid' column off of the DataFrame, just in case the user # has assigned it, and throw a warning try: df.pop('sid') warnings.warn( "Assignment of the 'sid' column of a DataFrame is " "not supported by Fetcher. The 'sid' column has been " "overwritten.", category=UserWarning, stacklevel=2, ) except KeyError: # There was no 'sid' column, so no warning is necessary pass # Fill entries for any symbols that don't require a date to # uniquely identify. Entries for which multiple securities exist # are replaced with zeroes, while entries for which no asset # exists are replaced with NaNs. unique_symbols = df[self.symbol_column].unique() sid_series = pd.Series( data=map(self._lookup_unconflicted_symbol, unique_symbols), index=unique_symbols, name='sid', ) df = df.join(sid_series, on=self.symbol_column) # Fill any zero entries left in our sid column by doing a lookup # using both symbol and the row date. conflict_rows = df[df['sid'] == 0] for row_idx, row in conflict_rows.iterrows(): try: asset = self.finder.lookup_symbol( row[self.symbol_column], # Replacing tzinfo here is necessary because of the # timezone metadata bug described below. row['dt'].replace(tzinfo=pytz.utc), # It's possible that no asset comes back here if our # lookup date is from before any asset held the # requested symbol. Mark such cases as NaN so that # they get dropped in the next step. ) or numpy.nan except SymbolNotFound: asset = numpy.nan # Assign the resolved asset to the cell df.ix[row_idx, 'sid'] = asset # Filter out rows containing symbols that we failed to find. length_before_drop = len(df) df = df[df['sid'].notnull()] no_sid_count = length_before_drop - len(df) if no_sid_count: logger.warn( "Dropped {} rows from fetched csv.".format(no_sid_count), no_sid_count, extra={'syslog': True}, ) else: df['sid'] = df['symbol'] # Dates are localized to UTC when they come out of # parse_date_str_series, but we need to re-localize them here because # of a bug that wasn't fixed until # https://github.com/pydata/pandas/pull/7092. # We should be able to remove the call to tz_localize once we're on # pandas 0.14.0 # We don't set 'dt' as the index until here because the Symbol parsing # operations above depend on having a unique index for the dataframe, # and the 'dt' column can contain multiple dates for the same entry. df.drop_duplicates(["sid", "dt"]) df.set_index(['dt'], inplace=True) df = df.tz_localize('UTC') df.sort_index(inplace=True) cols_to_drop = [self.date_column] if self.symbol is None: cols_to_drop.append(self.symbol_column) df = df[df.columns.drop(cols_to_drop)] if self.post_func: df = self.post_func(df) return df def __iter__(self): asset_cache = {} for dt, series in self.df.iterrows(): if dt < self.start_date: continue if dt > self.end_date: return event = FetcherEvent() # when dt column is converted to be the dataframe's index # the dt column is dropped. So, we need to manually copy # dt into the event. event.dt = dt for k, v in series.iteritems(): # convert numpy integer types to # int. This assumes we are on a 64bit # platform that will not lose information # by casting. # TODO: this is only necessary on the # amazon qexec instances. would be good # to figure out how to use the numpy dtypes # without this check and casting. if isinstance(v, numpy.integer): v = int(v) setattr(event, k, v) # If it has start_date, then it's already an Asset # object from asset_for_symbol, and we don't have to # transform it any further. Checking for start_date is # faster than isinstance. if event.sid in asset_cache: event.sid = asset_cache[event.sid] elif hasattr(event.sid, 'start_date'): # Clone for user algo code, if we haven't already. asset_cache[event.sid] = event.sid elif self.finder and isinstance(event.sid, int): asset = self.finder.retrieve_asset(event.sid, default_none=True) if asset: # Clone for user algo code. event.sid = asset_cache[asset] = asset elif self.mask: # When masking drop all non-mappable values. continue elif self.symbol is None: # If the event's sid property is an int we coerce # it into an Equity. event.sid = asset_cache[event.sid] = Equity(event.sid) event.type = DATASOURCE_TYPE.CUSTOM event.source_id = self.namestring yield event class PandasRequestsCSV(PandasCSV): # maximum 100 megs to prevent DDoS MAX_DOCUMENT_SIZE = (1024 * 1024) * 100 # maximum number of bytes to read in at a time CONTENT_CHUNK_SIZE = 4096 def __init__(self, url, pre_func, post_func, env, start_date, end_date, date_column, date_format, timezone, symbol, mask, symbol_column, data_frequency, special_params_checker=None, **kwargs): # Peel off extra requests kwargs, forwarding the remaining kwargs to # the superclass. # Also returns possible https updated url if sent to http quandl ds # If url hasn't changed, will just return the original. self._requests_kwargs, self.url =\ mask_requests_args(url, params_checker=special_params_checker, **kwargs) remaining_kwargs = { k: v for k, v in iteritems(kwargs) if k not in self.requests_kwargs } self.namestring = type(self).__name__ super(PandasRequestsCSV, self).__init__( pre_func, post_func, env, start_date, end_date, date_column, date_format, timezone, symbol, mask, symbol_column, data_frequency, **remaining_kwargs ) self.fetch_size = None self.fetch_hash = None self.df = self.load_df() self.special_params_checker = special_params_checker @property def requests_kwargs(self): return self._requests_kwargs def fetch_url(self, url): info = "checking {url} with {params}" logger.info(info.format(url=url, params=self.requests_kwargs)) # setting decode_unicode=True sometimes results in a # UnicodeEncodeError exception, so instead we'll use # pandas logic for decoding content try: response = requests.get(url, **self.requests_kwargs) except requests.exceptions.ConnectionError: raise Exception('Could not connect to %s' % url) if not response.ok: raise Exception('Problem reaching %s' % url) elif response.is_redirect: # On the offchance we don't catch a redirect URL # in validation, this will catch it. new_url = response.headers['location'] raise FetcherCSVRedirectError( url=url, new_url=new_url, extra={ 'old_url': url, 'new_url': new_url } ) content_length = 0 logger.info('{} connection established in {:.1f} seconds'.format( url, response.elapsed.total_seconds())) # use the decode_unicode flag to ensure that the output of this is # a string, and not bytes. for chunk in response.iter_content(self.CONTENT_CHUNK_SIZE, decode_unicode=True): if content_length > self.MAX_DOCUMENT_SIZE: raise Exception('Document size too big.') if chunk: content_length += len(chunk) yield chunk return def fetch_data(self): # create a data frame directly from the full text of # the response from the returned file-descriptor. data = self.fetch_url(self.url) fd = StringIO() if isinstance(data, str): fd.write(data) else: for chunk in data: fd.write(chunk) self.fetch_size = fd.tell() fd.seek(0) try: # see if pandas can parse csv data frames = read_csv(fd, **self.pandas_kwargs) frames_hash = hashlib.md5(str(fd.getvalue()).encode('utf-8')) self.fetch_hash = frames_hash.hexdigest() except pd.parser.CParserError: # could not parse the data, raise exception raise Exception('Error parsing remote CSV data.') finally: fd.close() return frames