Repository: quantopian/zipline
File: tests/pipeline/test_pipeline_algo.py (Python)

Code before:
    'resources',
    'pipeline_inputs',
)


def rolling_vwap(df, length):
    "Simple rolling vwap implementation for testing"
    closes = df['close'].values
    volumes = df['volume'].values
    product = closes * volumes
    out = full_like(closes, nan)
    for upper_bound in range(length, len(closes) + 1):
        bounds = slice(upper_bound - length, upper_bound)
        out[upper_bound - 1] = product[bounds].sum() / volumes[bounds].sum()

    return Series(out, index=df.index)


class ClosesAndVolumes(WithMakeAlgo, ZiplineTestCase):
    sids = 1, 2, 3
    START_DATE = pd.Timestamp('2014-01-01', tz='utc')
    END_DATE = pd.Timestamp('2014-02-01', tz='utc')
    dates = date_range(START_DATE, END_DATE, freq=get_calendar("NYSE").day,
                       tz='utc')

    SIM_PARAMS_DATA_FREQUENCY = 'daily'
    DATA_PORTAL_USE_MINUTE_DATA = False

    # FIXME: This currently uses benchmark returns from the trading
    # environment.
    BENCHMARK_SID = None

    @classmethod
    def make_equity_info(cls):
        cls.equity_info = ret = DataFrame.from_records([
            {
                'sid': 1,
                'symbol': 'A',
                'start_date': cls.dates[10],
                'end_date': cls.dates[13],
                'exchange': 'TEST',
            },
            {
                'sid': 2,
                'symbol': 'B',
                'start_date': cls.dates[11],
                'end_date': cls.dates[14],
                'exchange': 'TEST',
            },
            {
                'sid': 3,

Review comment (suggestion):
This is in service of a larger PR that aims to eventually kill `TradingEnvironment` entirely, so in the interest of changing as little as possible here, I'm not changing this behavior.

Code after:


def rolling_vwap(df, length):
    "Simple rolling vwap implementation for testing"
    closes = df['close'].values
    volumes = df['volume'].values
    product = closes * volumes
    out = full_like(closes, nan)
    for upper_bound in range(length, len(closes) + 1):
        bounds = slice(upper_bound - length, upper_bound)
        out[upper_bound - 1] = product[bounds].sum() / volumes[bounds].sum()

    return Series(out, index=df.index)


class ClosesAndVolumes(WithMakeAlgo, ZiplineTestCase):
    sids = 1, 2, 3
    START_DATE = pd.Timestamp('2014-01-01', tz='utc')
    END_DATE = pd.Timestamp('2014-02-01', tz='utc')
    dates = date_range(START_DATE, END_DATE, freq=get_calendar("NYSE").day,
                       tz='utc')

    SIM_PARAMS_DATA_FREQUENCY = 'daily'
    DATA_PORTAL_USE_MINUTE_DATA = False

    # FIXME: This currently uses benchmark returns from the trading
    # environment.
    BENCHMARK_SID = None

    @classmethod
    def make_equity_info(cls):
        cls.equity_info = ret = DataFrame.from_records([
            {
                'sid': 1,
                'symbol': 'A',
                'start_date': cls.dates[10],
                'end_date': cls.dates[13],
                'exchange': 'TEST',
            },
            {
                'sid': 2,
                'symbol': 'B',
                'start_date': cls.dates[11],
                'end_date': cls.dates[14],
                'exchange': 'TEST',
            },
            {
                'sid': 3,
                'symbol': 'C',
                'start_date': cls.dates[12],
                'end_date': cls.dates[15],