catalyst/catalyst/data/bundles/quandl.py

#
# Copyright 2017 Enigma MPC, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from datetime import datetime

import pandas as pd

from six.moves.urllib.parse import urlencode

from catalyst.data.bundles.core import register_bundle
from catalyst.data.bundles.base_pricing import BaseEquityPricingBundle
from catalyst.utils.memoize import lazyval

"""
Module for building a complete daily dataset from Quandl's WIKI dataset.
"""
from itertools import count
import tarfile
from time import time, sleep
from datetime import datetime

from logbook import Logger
import pandas as pd
from six.moves.urllib.parse import urlencode

from catalyst.utils.calendars import register_calendar_alias
from catalyst.utils.cli import maybe_show_progress

from . import core as bundles

from catalyst.constants import LOG_LEVEL

log = Logger(__name__, level=LOG_LEVEL)
seconds_per_call = (pd.Timedelta('10 minutes') / 2000).total_seconds()

class QuandlBundle(BaseEquityPricingBundle):
    @lazyval
    def name(self):
        return 'quandl'

    @lazyval
    def exchange(self):
        return 'QUANDL'

    @lazyval
    def frequencies(self):
        return set(('daily',))

    @lazyval
    def tar_url(self):
        return 'https://s3.amazonaws.com/quantopian-public-zipline-data/quandl'

    @lazyval
    def wait_time(self):
        return pd.Timedelta(milliseconds=300)

    @lazyval
    def _excluded_symbols(self):
        """
        Invalid symbols that quandl has had in its metadata:
        """
        return frozenset({'TEST123456789'})

    def fetch_raw_metadata_frame(self, api_key, page_number):
        raw = pd.read_csv(
            self._format_metadata_url(api_key, page_number),
            date_parser=pd.tseries.tools.to_datetime,
            parse_dates=[
                'oldest_available_date',
                'newest_available_date',
            ],
            dtype={
                'dataset_code': 'str',
                'name': 'str',
                'oldest_available_date': 'str',
                'newest_available_date': 'str',
            },
            usecols=[
                'dataset_code',
                'name',
                'oldest_available_date',
                'newest_available_date',
            ],
        ).rename(
            columns={
                'dataset_code': 'symbol',
                'name': 'asset_name',
                'oldest_available_date': 'start_date',
                'newest_available_date': 'end_date',
            },
        )

        raw['start_date'] = raw['start_date'].astype(datetime)
        raw['end_date'] = raw['end_date'].astype(datetime)
        raw['ac_date'] = raw['end_date'] + pd.Timedelta(days=1)

        # Filter out invalid symbols
        raw = raw[~raw.symbol.isin(self._excluded_symbols)]

        # cut out all the other stuff in the name column
        # we need to escape the paren because it is actually splitting on a regex
        raw.asset_name = raw.asset_name.str.split(r' \(', 1).str.get(0)

        return raw

    def fetch_raw_symbol_frame(self,
                               api_key,
                               symbol,
                               calendar,
                               start_session,
                               end_session,
                               data_frequency):
        raw_data = pd.read_csv(
            self._format_wiki_url(
                api_key,
                symbol,
                start_session,
                end_session,
                data_frequency,
            ),
            parse_dates=['Date'],
            index_col='Date',
            usecols=[
                'Open',
                'High',
                'Low',
                'Close',
                'Volume',
                'Date',
                'Ex-Dividend',
                'Split Ratio',
            ],
            na_values=['NA'],
        ).rename(columns={
            'Open': 'open',
            'High': 'high',
            'Low': 'low',
            'Close': 'close',
            'Volume': 'volume',
            'Date': 'date',
            'Ex-Dividend': 'ex_dividend',
            'Split Ratio': 'split_ratio',
        })

        sessions = calendar.sessions_in_range(start_session, end_session)

        return raw_data.reindex(
            sessions.tz_localize(None),
            copy=False,
        ).fillna(0.0)

    def post_process_symbol_metadata(self, asset_id, sym_md, sym_data):
        self._update_splits(asset_id, sym_data)
        self._update_dividends(asset_id, sym_data)

        return sym_md

    def _update_splits(self, asset_id, raw_data):
        split_ratios = raw_data.split_ratio
        df = pd.DataFrame({'ratio': 1 / split_ratios[split_ratios != 1]})
        df.index.name = 'effective_date'
        df.reset_index(inplace=True)
        df['sid'] = asset_id
        self.splits.append(df)


    def _update_dividends(self, asset_id, raw_data):
        divs = raw_data.ex_dividend
        df = pd.DataFrame({'amount': divs[divs != 0]})
        df.index.name = 'ex_date'
        df.reset_index(inplace=True)
        df['sid'] = asset_id
        # we do not have this data in the WIKI dataset
        df['record_date'] = df['declared_date'] = df['pay_date'] = pd.NaT
        self.dividends.append(df)


    def _format_metadata_url(self, api_key, page_number):
        """Build the query RL for the quandl WIKI metadata.
        """
        query_params = [
            ('per_page', '100'),
            ('sort_by', 'id'),
            ('page', str(page_number)),
            ('database_code', 'WIKI'),
        ]
        if api_key is not None:
            query_params = [('api_key', api_key)] + query_params

        return (
            'https://www.quandl.com/api/v3/datasets.csv?' + urlencode(query_params)
        )


    def _format_wiki_url(self,
                         api_key,
                         symbol,
                         start_date,
                         end_date,
                         data_frequency):
        """
        Build a query URL for a quandl WIKI dataset.
        """
        query_params = [
            ('start_date', start_date.strftime('%Y-%m-%d')),
            ('end_date', end_date.strftime('%Y-%m-%d')),
            ('order', 'asc'),
        ]
        if api_key is not None:
            query_params = [('api_key', api_key)] + query_params

        return (
            "https://www.quandl.com/api/v3/datasets/WIKI/"
            "{symbol}.csv?{query}".format(
                symbol=symbol,
                query=urlencode(query_params),
            )
        )

register_calendar_alias('QUANDL', 'NYSE')
register_bundle(QuandlBundle)