<s> from datetime import tzinfo <EOL> from functools import partial , wraps <EOL> from operator import attrgetter <EOL> from numpy import dtype <EOL> from pytz import timezone <EOL> from six import iteritems , string_types , PY3 <EOL> from toolz import valmap , complement , compose <EOL> import toolz . curried . operator as op <EOL> from zipline . utils . preprocess import preprocess <EOL> def optionally ( preprocessor ) : <EOL> """<STR_LIT>""" <EOL> @ wraps ( preprocessor ) <EOL> def wrapper ( func , argname , arg ) : <EOL> return arg if arg is None else preprocessor ( func , argname , arg ) <EOL> return wrapper <EOL> def ensure_upper_case ( func , argname , arg ) : <EOL> if isinstance ( arg , string_types ) : <EOL> return arg . upper ( ) <EOL> else : <EOL> raise TypeError ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" . format ( <EOL> func . __name__ , <EOL> argname , <EOL> arg , <EOL> ) , <EOL> ) <EOL> def ensure_dtype ( func , argname , arg ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> return dtype ( arg ) <EOL> except TypeError : <EOL> raise TypeError ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" . format ( <EOL> func = _qualified_name ( func ) , <EOL> argname = argname , <EOL> arg = arg , <EOL> ) , <EOL> ) <EOL> def ensure_timezone ( func , argname , arg ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( arg , tzinfo ) : <EOL> return arg <EOL> if isinstance ( arg , string_types ) : <EOL> return timezone ( arg ) <EOL> raise TypeError ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" . format ( <EOL> func = _qualified_name ( func ) , <EOL> argname = argname , <EOL> arg = arg , <EOL> ) , <EOL> ) <EOL> def expect_dtypes ( * _pos , ** named ) : <EOL> """<STR_LIT>""" <EOL> if _pos : <EOL> raise TypeError ( "<STR_LIT>" ) <EOL> for name , type_ in iteritems ( named ) : <EOL> if not isinstance ( type_ , ( dtype , tuple ) ) : <EOL> raise TypeError ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" . format ( <EOL> name = name , dtype = dtype , <EOL> ) <EOL> ) <EOL> def _expect_dtype ( _dtype_or_dtype_tuple ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( _dtype_or_dtype_tuple , tuple ) : <EOL> allowed_dtypes = _dtype_or_dtype_tuple <EOL> else : <EOL> allowed_dtypes = ( _dtype_or_dtype_tuple , ) <EOL> template = ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> ) . format ( dtype_str = '<STR_LIT>' . join ( repr ( d . name ) for d in allowed_dtypes ) ) <EOL> def check_dtype ( value ) : <EOL> return getattr ( value , '<STR_LIT>' , None ) not in allowed_dtypes <EOL> def display_bad_value ( value ) : <EOL> try : <EOL> return value . dtype . name <EOL> except AttributeError : <EOL> return value <EOL> return make_check ( <EOL> exc_type = TypeError , <EOL> template = template , <EOL> pred = check_dtype , <EOL> actual = display_bad_value , <EOL> ) <EOL> return preprocess ( ** valmap ( _expect_dtype , named ) ) <EOL> def expect_types ( * _pos , ** named ) : <EOL> """<STR_LIT>""" <EOL> if _pos : <EOL> raise TypeError ( "<STR_LIT>" ) <EOL> for name , type_ in iteritems ( named ) : <EOL> if not isinstance ( type_ , ( type , tuple ) ) : <EOL> raise TypeError ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" . format ( <EOL> name = name , type_ = type_ , <EOL> ) <EOL> ) <EOL> def _expect_type ( type_ ) : <EOL> _template = ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> ) <EOL> if isinstance ( type_ , tuple ) : <EOL> template = _template . format ( <EOL> type_or_types = '<STR_LIT>' . join ( map ( _qualified_name , type_ ) ) <EOL> ) <EOL> else : <EOL> template = _template . format ( type_or_types = _qualified_name ( type_ ) ) <EOL> return make_check ( <EOL> TypeError , <EOL> template , <EOL> lambda v : not isinstance ( v , type_ ) , <EOL> compose ( _qualified_name , type ) , <EOL> ) <EOL> return preprocess ( ** valmap ( _expect_type , named ) ) <EOL> if PY3 : <EOL> _qualified_name = attrgetter ( '<STR_LIT>' ) <EOL> else : <EOL> def _qualified_name ( obj ) : <EOL> """<STR_LIT>""" <EOL> module = obj . __module__ <EOL> if module in ( '<STR_LIT>' , '<STR_LIT:__main__>' , '<STR_LIT>' ) : <EOL> return obj . __name__ <EOL> return '<STR_LIT:.>' . join ( [ module , obj . __name__ ] ) <EOL> def make_check ( exc_type , template , pred , actual ) : <EOL> """<STR_LIT>""" <EOL> def _check ( func , argname , argvalue ) : <EOL> if pred ( argvalue ) : <EOL> raise exc_type ( <EOL> template % { <EOL> '<STR_LIT>' : _qualified_name ( func ) , <EOL> '<STR_LIT>' : argname , <EOL> '<STR_LIT>' : actual ( argvalue ) , <EOL> } , <EOL> ) <EOL> return argvalue <EOL> return _check <EOL> def optional ( type_ ) : <EOL> """<STR_LIT>""" <EOL> return ( type_ , type ( None ) ) <EOL> def expect_element ( * _pos , ** named ) : <EOL> """<STR_LIT>""" <EOL> if _pos : <EOL> raise TypeError ( "<STR_LIT>" ) <EOL> def _expect_element ( collection ) : <EOL> template = ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> ) . format ( collection = collection ) <EOL> return make_check ( <EOL> ValueError , <EOL> template , <EOL> complement ( op . contains ( collection ) ) , <EOL> repr , <EOL> ) <EOL> return preprocess ( ** valmap ( _expect_element , named ) ) <EOL> def expect_dimensions ( ** dimensions ) : <EOL> """<STR_LIT>""" <EOL> def _expect_dimension ( expected_ndim ) : <EOL> def _check ( func , argname , argvalue ) : <EOL> funcname = _qualified_name ( func ) <EOL> actual_ndim = argvalue . ndim <EOL> if actual_ndim != expected_ndim : <EOL> if actual_ndim == <NUM_LIT:0> : <EOL> actual_repr = '<STR_LIT>' <EOL> else : <EOL> actual_repr = "<STR_LIT>" % actual_ndim <EOL> raise ValueError ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" . format ( <EOL> func = funcname , <EOL> expected = expected_ndim , <EOL> argname = argname , <EOL> actual = actual_repr , <EOL> ) <EOL> ) <EOL> return argvalue <EOL> return _check <EOL> return preprocess ( ** valmap ( _expect_dimension , dimensions ) ) <EOL> def coerce ( from_ , to , ** to_kwargs ) : <EOL> """<STR_LIT>""" <EOL> def preprocessor ( func , argname , arg ) : <EOL> if isinstance ( arg , from_ ) : <EOL> return to ( arg , ** to_kwargs ) <EOL> return arg <EOL> return preprocessor <EOL> class error_keywords ( object ) : <EOL> def __init__ ( self , * args , ** kwargs ) : <EOL> self . messages = kwargs <EOL> def __call__ ( self , func ) : <EOL> def assert_keywords_and_call ( * args , ** kwargs ) : <EOL> for field , message in iteritems ( self . messages ) : <EOL> if field in kwargs : <EOL> raise TypeError ( message ) <EOL> return func ( * args , ** kwargs ) <EOL> return assert_keywords_and_call <EOL> coerce_string = partial ( coerce , string_types ) </s>
