<s> from six import StringIO , iteritems <EOL> from abc import ABCMeta , abstractmethod <EOL> from collections import namedtuple <EOL> import hashlib <EOL> from textwrap import dedent <EOL> import pandas as pd <EOL> from pandas import read_csv <EOL> import numpy <EOL> from logbook import Logger <EOL> import pytz <EOL> import warnings <EOL> import requests <EOL> from zipline . errors import ( <EOL> MultipleSymbolsFound , <EOL> SymbolNotFound , <EOL> ZiplineError <EOL> ) <EOL> from zipline . protocol import ( <EOL> DATASOURCE_TYPE , <EOL> Event <EOL> ) <EOL> from zipline . assets import Equity <EOL> logger = Logger ( '<STR_LIT>' ) <EOL> def roll_dts_to_midnight ( dts , env ) : <EOL> if len ( dts ) == <NUM_LIT:0> : <EOL> return dts <EOL> return pd . DatetimeIndex ( <EOL> ( dts . tz_convert ( '<STR_LIT>' ) - pd . Timedelta ( hours = <NUM_LIT:16> ) ) . date , <EOL> tz = '<STR_LIT>' , <EOL> ) + env . trading_day <EOL> class FetcherEvent ( Event ) : <EOL> pass <EOL> class FetcherCSVRedirectError ( ZiplineError ) : <EOL> msg = dedent ( <EOL> """<STR_LIT>""" <EOL> ) <EOL> def __init__ ( self , * args , ** kwargs ) : <EOL> self . url = kwargs [ "<STR_LIT:url>" ] <EOL> self . new_url = kwargs [ "<STR_LIT>" ] <EOL> self . extra = kwargs [ "<STR_LIT>" ] <EOL> super ( FetcherCSVRedirectError , self ) . __init__ ( * args , ** kwargs ) <EOL> ALLOWED_REQUESTS_KWARGS = { <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' } <EOL> ALLOWED_READ_CSV_KWARGS = { <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' <EOL> } <EOL> SHARED_REQUESTS_KWARGS = { <EOL> '<STR_LIT>' : True , <EOL> '<STR_LIT>' : False , <EOL> } <EOL> def mask_requests_args ( url , validating = False , params_checker = None , ** kwargs ) : <EOL> requests_kwargs = { key : val for ( key , val ) in iteritems ( kwargs ) <EOL> if key in ALLOWED_REQUESTS_KWARGS } <EOL> if params_checker is not None : <EOL> url , s_params = params_checker ( url ) <EOL> if s_params : <EOL> if '<STR_LIT>' in requests_kwargs : <EOL> requests_kwargs [ '<STR_LIT>' ] . update ( s_params ) <EOL> else : <EOL> requests_kwargs [ '<STR_LIT>' ] = s_params <EOL> requests_kwargs [ '<STR_LIT>' ] = <NUM_LIT:1.0> if validating else <NUM_LIT> <EOL> requests_kwargs . update ( SHARED_REQUESTS_KWARGS ) <EOL> request_pair = namedtuple ( "<STR_LIT>" , ( "<STR_LIT>" , "<STR_LIT:url>" ) ) <EOL> return request_pair ( requests_kwargs , url ) <EOL> class PandasCSV ( object ) : <EOL> __metaclass__ = ABCMeta <EOL> def __init__ ( self , <EOL> pre_func , <EOL> post_func , <EOL> env , <EOL> start_date , <EOL> end_date , <EOL> date_column , <EOL> date_format , <EOL> timezone , <EOL> symbol , <EOL> mask , <EOL> symbol_column , <EOL> data_frequency , <EOL> ** kwargs ) : <EOL> self . start_date = start_date <EOL> self . end_date = end_date <EOL> self . date_column = date_column <EOL> self . date_format = date_format <EOL> self . timezone = timezone <EOL> self . mask = mask <EOL> self . symbol_column = symbol_column or "<STR_LIT>" <EOL> self . data_frequency = data_frequency <EOL> invalid_kwargs = set ( kwargs ) - ALLOWED_READ_CSV_KWARGS <EOL> if invalid_kwargs : <EOL> raise TypeError ( <EOL> "<STR_LIT>" % invalid_kwargs , <EOL> ) <EOL> self . pandas_kwargs = self . mask_pandas_args ( kwargs ) <EOL> self . symbol = symbol <EOL> self . env = env <EOL> self . finder = env . asset_finder <EOL> self . pre_func = pre_func <EOL> self . post_func = post_func <EOL> @ property <EOL> def fields ( self ) : <EOL> return self . df . columns . tolist ( ) <EOL> def get_hash ( self ) : <EOL> return self . namestring <EOL> @ abstractmethod <EOL> def fetch_data ( self ) : <EOL> return <EOL> @ staticmethod <EOL> def parse_date_str_series ( format_str , tz , date_str_series , data_frequency , <EOL> env ) : <EOL> """<STR_LIT>""" <EOL> if format_str is not None : <EOL> logger . warn ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> ) <EOL> format_str = None <EOL> tz_str = str ( tz ) <EOL> if tz_str == pytz . utc . zone : <EOL> parsed = pd . to_datetime ( <EOL> date_str_series . values , <EOL> format = format_str , <EOL> utc = True , <EOL> coerce = True , <EOL> ) <EOL> else : <EOL> parsed = pd . to_datetime ( <EOL> date_str_series . values , <EOL> format = format_str , <EOL> coerce = True , <EOL> ) . tz_localize ( tz_str ) . tz_convert ( '<STR_LIT>' ) <EOL> if data_frequency == '<STR_LIT>' : <EOL> parsed = roll_dts_to_midnight ( parsed , env ) <EOL> return parsed <EOL> def mask_pandas_args ( self , kwargs ) : <EOL> pandas_kwargs = { key : val for ( key , val ) in iteritems ( kwargs ) <EOL> if key in ALLOWED_READ_CSV_KWARGS } <EOL> if '<STR_LIT>' in pandas_kwargs : <EOL> usecols = pandas_kwargs [ '<STR_LIT>' ] <EOL> if usecols and self . date_column not in usecols : <EOL> with_date = list ( usecols ) <EOL> with_date . append ( self . date_column ) <EOL> pandas_kwargs [ '<STR_LIT>' ] = with_date <EOL> pandas_kwargs . setdefault ( '<STR_LIT>' , False ) <EOL> pandas_kwargs . setdefault ( '<STR_LIT>' , { '<STR_LIT>' : [ ] } ) <EOL> return pandas_kwargs <EOL> def _lookup_unconflicted_symbol ( self , symbol ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> uppered = symbol . upper ( ) <EOL> except AttributeError : <EOL> return numpy . nan <EOL> try : <EOL> return self . finder . lookup_symbol ( uppered , as_of_date = None ) <EOL> except MultipleSymbolsFound : <EOL> return <NUM_LIT:0> <EOL> except SymbolNotFound : <EOL> return numpy . nan <EOL> def load_df ( self ) : <EOL> df = self . fetch_data ( ) <EOL> if self . pre_func : <EOL> df = self . pre_func ( df ) <EOL> df [ '<STR_LIT>' ] = self . parse_date_str_series ( <EOL> self . date_format , <EOL> self . timezone , <EOL> df [ self . date_column ] , <EOL> self . data_frequency , <EOL> self . env <EOL> ) . values <EOL> df = df [ df [ '<STR_LIT>' ] . notnull ( ) ] <EOL> if self . symbol is not None : <EOL> df [ '<STR_LIT>' ] = self . symbol <EOL> elif self . finder : <EOL> df . sort ( self . symbol_column ) <EOL> try : <EOL> df . pop ( '<STR_LIT>' ) <EOL> warnings . warn ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" , <EOL> category = UserWarning , <EOL> stacklevel = <NUM_LIT:2> , <EOL> ) <EOL> except KeyError : <EOL> pass <EOL> unique_symbols = df [ self . symbol_column ] . unique ( ) <EOL> sid_series = pd . Series ( <EOL> data = map ( self . _lookup_unconflicted_symbol , unique_symbols ) , <EOL> index = unique_symbols , <EOL> name = '<STR_LIT>' , <EOL> ) <EOL> df = df . join ( sid_series , on = self . symbol_column ) <EOL> conflict_rows = df [ df [ '<STR_LIT>' ] == <NUM_LIT:0> ] <EOL> for row_idx , row in conflict_rows . iterrows ( ) : <EOL> try : <EOL> asset = self . finder . lookup_symbol ( <EOL> row [ self . symbol_column ] , <EOL> row [ '<STR_LIT>' ] . replace ( tzinfo = pytz . utc ) , <EOL> ) or numpy . nan <EOL> except SymbolNotFound : <EOL> asset = numpy . nan <EOL> df . ix [ row_idx , '<STR_LIT>' ] = asset <EOL> length_before_drop = len ( df ) <EOL> df = df [ df [ '<STR_LIT>' ] . notnull ( ) ] <EOL> no_sid_count = length_before_drop - len ( df ) <EOL> if no_sid_count : <EOL> logger . warn ( <EOL> "<STR_LIT>" . format ( no_sid_count ) , <EOL> no_sid_count , <EOL> extra = { '<STR_LIT>' : True } , <EOL> ) <EOL> else : <EOL> df [ '<STR_LIT>' ] = df [ '<STR_LIT>' ] <EOL> df . drop_duplicates ( [ "<STR_LIT>" , "<STR_LIT>" ] ) <EOL> df . set_index ( [ '<STR_LIT>' ] , inplace = True ) <EOL> df = df . tz_localize ( '<STR_LIT>' ) <EOL> df . sort_index ( inplace = True ) <EOL> cols_to_drop = [ self . date_column ] <EOL> if self . symbol is None : <EOL> cols_to_drop . append ( self . symbol_column ) <EOL> df = df [ df . columns . drop ( cols_to_drop ) ] <EOL> if self . post_func : <EOL> df = self . post_func ( df ) <EOL> return df <EOL> def __iter__ ( self ) : <EOL> asset_cache = { } <EOL> for dt , series in self . df . iterrows ( ) : <EOL> if dt < self . start_date : <EOL> continue <EOL> if dt > self . end_date : <EOL> return <EOL> event = FetcherEvent ( ) <EOL> event . dt = dt <EOL> for k , v in series . iteritems ( ) : <EOL> if isinstance ( v , numpy . integer ) : <EOL> v = int ( v ) <EOL> setattr ( event , k , v ) <EOL> if event . sid in asset_cache : <EOL> event . sid = asset_cache [ event . sid ] <EOL> elif hasattr ( event . sid , '<STR_LIT>' ) : <EOL> asset_cache [ event . sid ] = event . sid <EOL> elif self . finder and isinstance ( event . sid , int ) : <EOL> asset = self . finder . retrieve_asset ( event . sid , <EOL> default_none = True ) <EOL> if asset : <EOL> event . sid = asset_cache [ asset ] = asset <EOL> elif self . mask : <EOL> continue <EOL> elif self . symbol is None : <EOL> event . sid = asset_cache [ event . sid ] = Equity ( event . sid ) <EOL> event . type = DATASOURCE_TYPE . CUSTOM <EOL> event . source_id = self . namestring <EOL> yield event <EOL> class PandasRequestsCSV ( PandasCSV ) : <EOL> MAX_DOCUMENT_SIZE = ( <NUM_LIT> * <NUM_LIT> ) * <NUM_LIT:100> <EOL> CONTENT_CHUNK_SIZE = <NUM_LIT> <EOL> def __init__ ( self , <EOL> url , <EOL> pre_func , <EOL> post_func , <EOL> env , <EOL> start_date , <EOL> end_date , <EOL> date_column , <EOL> date_format , <EOL> timezone , <EOL> symbol , <EOL> mask , <EOL> symbol_column , <EOL> data_frequency , <EOL> special_params_checker = None , <EOL> ** kwargs ) : <EOL> self . _requests_kwargs , self . url = mask_requests_args ( url , <EOL> params_checker = special_params_checker , <EOL> ** kwargs ) <EOL> remaining_kwargs = { <EOL> k : v for k , v in iteritems ( kwargs ) <EOL> if k not in self . requests_kwargs <EOL> } <EOL> self . namestring = type ( self ) . __name__ <EOL> super ( PandasRequestsCSV , self ) . __init__ ( <EOL> pre_func , <EOL> post_func , <EOL> env , <EOL> start_date , <EOL> end_date , <EOL> date_column , <EOL> date_format , <EOL> timezone , <EOL> symbol , <EOL> mask , <EOL> symbol_column , <EOL> data_frequency , <EOL> ** remaining_kwargs <EOL> ) <EOL> self . fetch_size = None <EOL> self . fetch_hash = None <EOL> self . df = self . load_df ( ) <EOL> self . special_params_checker = special_params_checker <EOL> @ property <EOL> def requests_kwargs ( self ) : <EOL> return self . _requests_kwargs <EOL> def fetch_url ( self , url ) : <EOL> info = "<STR_LIT>" <EOL> logger . info ( info . format ( url = url , params = self . requests_kwargs ) ) <EOL> try : <EOL> response = requests . get ( url , ** self . requests_kwargs ) <EOL> except requests . exceptions . ConnectionError : <EOL> raise Exception ( '<STR_LIT>' % url ) <EOL> if not response . ok : <EOL> raise Exception ( '<STR_LIT>' % url ) <EOL> elif response . is_redirect : <EOL> new_url = response . headers [ '<STR_LIT:location>' ] <EOL> raise FetcherCSVRedirectError ( <EOL> url = url , <EOL> new_url = new_url , <EOL> extra = { <EOL> '<STR_LIT>' : url , <EOL> '<STR_LIT>' : new_url <EOL> } <EOL> ) <EOL> content_length = <NUM_LIT:0> <EOL> logger . info ( '<STR_LIT>' . format ( <EOL> url , response . elapsed . total_seconds ( ) ) ) <EOL> for chunk in response . iter_content ( self . CONTENT_CHUNK_SIZE , <EOL> decode_unicode = True ) : <EOL> if content_length > self . MAX_DOCUMENT_SIZE : <EOL> raise Exception ( '<STR_LIT>' ) <EOL> if chunk : <EOL> content_length += len ( chunk ) <EOL> yield chunk <EOL> return <EOL> def fetch_data ( self ) : <EOL> data = self . fetch_url ( self . url ) <EOL> fd = StringIO ( ) <EOL> if isinstance ( data , str ) : <EOL> fd . write ( data ) <EOL> else : <EOL> for chunk in data : <EOL> fd . write ( chunk ) <EOL> self . fetch_size = fd . tell ( ) <EOL> fd . seek ( <NUM_LIT:0> ) <EOL> try : <EOL> frames = read_csv ( fd , ** self . pandas_kwargs ) <EOL> frames_hash = hashlib . md5 ( str ( fd . getvalue ( ) ) . encode ( '<STR_LIT:utf-8>' ) ) <EOL> self . fetch_hash = frames_hash . hexdigest ( ) <EOL> except pd . parser . CParserError : <EOL> raise Exception ( '<STR_LIT>' ) <EOL> finally : <EOL> fd . close ( ) <EOL> return frames </s>
