<s> from functools import wraps <EOL> from statsmodels . tools . data import _is_using_pandas <EOL> from statsmodels . tsa . base import datetools <EOL> from statsmodels . tsa . tsatools import freq_to_period <EOL> def _get_pandas_wrapper ( X , trim_head = None , trim_tail = None , names = None ) : <EOL> index = X . index <EOL> if trim_head is None and trim_tail is None : <EOL> index = index <EOL> elif trim_tail is None : <EOL> index = index [ trim_head : ] <EOL> elif trim_head is None : <EOL> index = index [ : - trim_tail ] <EOL> else : <EOL> index = index [ trim_head : - trim_tail ] <EOL> if hasattr ( X , "<STR_LIT>" ) : <EOL> if names is None : <EOL> names = X . columns <EOL> return lambda x : X . __class__ ( x , index = index , columns = names ) <EOL> else : <EOL> if names is None : <EOL> names = X . name <EOL> return lambda x : X . __class__ ( x , index = index , name = names ) <EOL> def _maybe_get_pandas_wrapper ( X , trim_head = None , trim_tail = None ) : <EOL> """<STR_LIT>""" <EOL> if _is_using_pandas ( X , None ) : <EOL> return _get_pandas_wrapper ( X , trim_head , trim_tail ) <EOL> else : <EOL> return <EOL> def _maybe_get_pandas_wrapper_freq ( X , trim = None ) : <EOL> if _is_using_pandas ( X , None ) : <EOL> index = X . index <EOL> func = _get_pandas_wrapper ( X , trim ) <EOL> freq = index . inferred_freq <EOL> return func , freq <EOL> else : <EOL> return lambda x : x , None <EOL> def pandas_wrapper ( func , trim_head = None , trim_tail = None , names = None , * args , <EOL> ** kwargs ) : <EOL> @ wraps ( func ) <EOL> def new_func ( X , * args , ** kwargs ) : <EOL> if not _is_using_pandas ( X , None ) : <EOL> return func ( X , * args , ** kwargs ) <EOL> wrapper_func = _get_pandas_wrapper ( X , trim_head , trim_tail , <EOL> names ) <EOL> ret = func ( X , * args , ** kwargs ) <EOL> ret = wrapper_func ( ret ) <EOL> return ret <EOL> return new_func <EOL> def pandas_wrapper_bunch ( func , trim_head = None , trim_tail = None , <EOL> names = None , * args , ** kwargs ) : <EOL> @ wraps ( func ) <EOL> def new_func ( X , * args , ** kwargs ) : <EOL> if not _is_using_pandas ( X , None ) : <EOL> return func ( X , * args , ** kwargs ) <EOL> wrapper_func = _get_pandas_wrapper ( X , trim_head , trim_tail , <EOL> names ) <EOL> ret = func ( X , * args , ** kwargs ) <EOL> ret = wrapper_func ( ret ) <EOL> return ret <EOL> return new_func <EOL> def pandas_wrapper_predict ( func , trim_head = None , trim_tail = None , <EOL> columns = None , * args , ** kwargs ) : <EOL> pass <EOL> def pandas_wrapper_freq ( func , trim_head = None , trim_tail = None , <EOL> freq_kw = '<STR_LIT>' , columns = None , * args , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> @ wraps ( func ) <EOL> def new_func ( X , * args , ** kwargs ) : <EOL> if not _is_using_pandas ( X , None ) : <EOL> return func ( X , * args , ** kwargs ) <EOL> wrapper_func = _get_pandas_wrapper ( X , trim_head , trim_tail , <EOL> columns ) <EOL> index = X . index <EOL> freq = index . inferred_freq <EOL> kwargs . update ( { freq_kw : freq_to_period ( freq ) } ) <EOL> ret = func ( X , * args , ** kwargs ) <EOL> ret = wrapper_func ( ret ) <EOL> return ret <EOL> return new_func <EOL> def dummy_func ( X ) : <EOL> return X <EOL> def dummy_func_array ( X ) : <EOL> return X . values <EOL> def dummy_func_pandas_columns ( X ) : <EOL> return X . values <EOL> def dummy_func_pandas_series ( X ) : <EOL> return X [ '<STR_LIT:A>' ] <EOL> import pandas as pd <EOL> import numpy as np <EOL> def test_pandas_freq_decorator ( ) : <EOL> X = pd . util . testing . makeDataFrame ( ) <EOL> func = pandas_wrapper ( dummy_func ) <EOL> np . testing . assert_equal ( func ( X . values ) , X ) <EOL> func = pandas_wrapper ( dummy_func_array ) <EOL> pd . util . testing . assert_frame_equal ( func ( X ) , X ) <EOL> expected = X . rename ( columns = dict ( zip ( '<STR_LIT>' , '<STR_LIT>' ) ) ) <EOL> func = pandas_wrapper ( dummy_func_array , names = list ( '<STR_LIT>' ) ) <EOL> pd . util . testing . assert_frame_equal ( func ( X ) , expected ) </s>
