zipline.sources.PandasRequestsCSV

Complexity

Total Complexity

15

Size/Duplication

Total Lines	138
Duplicated Lines	0 %

4 Methods

Rating	Name	Duplication	Size	Complexity
A	__init__()	0	55	3
A	requests_kwargs()	0	3	1
B	fetch_data()	0	29	4
C	fetch_url()	0	41	7

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):
    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