<commit_before>#
# Copyright 2013 Quantopian, 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.


"""
Factory functions to prepare useful data.
"""
import pytz
import random
from collections import OrderedDict
from delorean import Delorean

import pandas as pd
import numpy as np
from datetime import datetime, timedelta

from zipline.protocol import Event, DATASOURCE_TYPE
from zipline.sources import (SpecificEquityTrades,
                             DataFrameSource,
                             DataPanelSource)
from zipline.finance.trading import SimulationParameters
from zipline.finance import trading
from zipline.sources.test_source import (
    date_gen,
    create_trade
)

# For backwards compatibility
from zipline.data.loader import (load_from_yahoo,
                                 load_bars_from_yahoo)

__all__ = ['load_from_yahoo', 'load_bars_from_yahoo']


def create_simulation_parameters(year=2006, start=None, end=None,
                                 capital_base=float("1.0e5"),
                                 num_days=None, load=None
                                 ):
    """Construct a complete environment with reasonable defaults"""
    if start is None:
        start = datetime(year, 1, 1, tzinfo=pytz.utc)
    if end is None:
        if num_days:
            trading.environment = trading.TradingEnvironment(load=load)
            start_index = trading.environment.trading_days.searchsorted(
                start)
            end = trading.environment.trading_days[start_index + num_days - 1]
        else:
            end = datetime(year, 12, 31, tzinfo=pytz.utc)
    sim_params = SimulationParameters(
        period_start=start,
        period_end=end,
        capital_base=capital_base,
    )

    return sim_params


def create_noop_environment():
    oneday = timedelta(days=1)
    start = datetime(2006, 1, 1, tzinfo=pytz.utc)

    days = []
    tr_curves = OrderedDict()
    for day in date_gen(start=start, delta=oneday, count=252):
        days.append(day)
        curve = {
            '10year': 0.0799,
            '1month': 0.0799,
            '1year': 0.0785,
            '20year': 0.0765,
            '2year': 0.0794,
            '30year': 0.0804,
            '3month': 0.0789,
            '3year': 0.0796,
            '5year': 0.0792,
            '6month': 0.0794,
            '7year': 0.0804,
            'tid': 1752
        }
        tr_curves[day] = curve

    bm_returns = pd.Series(index=days, data=0.1)

    load_nodata = lambda x: (bm_returns, tr_curves)

    return trading.TradingEnvironment(load=load_nodata)


def create_random_simulation_parameters():
    trading.environment = trading.TradingEnvironment()
    treasury_curves = trading.environment.treasury_curves

    for n in range(100):

        random_index = random.randint(
            0,
            len(treasury_curves) - 1
        )

        start_dt = treasury_curves.index[random_index]
        end_dt = start_dt + timedelta(days=365)

        now = datetime.utcnow().replace(tzinfo=pytz.utc)

        if end_dt <= now:
            break

    assert end_dt <= now, """
failed to find a suitable daterange after 100 attempts. please double
check treasury and benchmark data in findb, and re-run the test."""

    sim_params = SimulationParameters(
        period_start=start_dt,
        period_end=end_dt
    )

    return sim_params, start_dt, end_dt


def get_next_trading_dt(current, interval):
    naive = current.replace(tzinfo=None)
    delo = Delorean(naive, pytz.utc.zone)
    ex_tz = trading.environment.exchange_tz
    next_dt = delo.shift(ex_tz).datetime

    while True:
        next_dt = next_dt + interval
        next_delo = Delorean(next_dt.replace(tzinfo=None), ex_tz)
        next_utc = next_delo.shift(pytz.utc.zone).datetime
        if trading.environment.is_market_hours(next_utc):
            break

    return next_utc


def create_trade_history(sid, prices, amounts, interval, sim_params,
                         source_id="test_factory"):
    trades = []
    current = sim_params.first_open

    oneday = timedelta(days=1)
    use_midnight = interval >= oneday
    for price, amount in zip(prices, amounts):
        if use_midnight:
            trade_dt = current.replace(hour=0, minute=0)
        else:
            trade_dt = current
        trade = create_trade(sid, price, amount, trade_dt, source_id)
        trades.append(trade)
        current = get_next_trading_dt(current, interval)

    assert len(trades) == len(prices)
    return trades<commit_msg>MAINT: Use Timestamp for calculating next_trading_dt in test factory.<commit_after>#
# Copyright 2013 Quantopian, 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.


"""
Factory functions to prepare useful data.
"""
import pytz
import random
from collections import OrderedDict

import pandas as pd
import numpy as np
from datetime import datetime, timedelta

from zipline.protocol import Event, DATASOURCE_TYPE
from zipline.sources import (SpecificEquityTrades,
                             DataFrameSource,
                             DataPanelSource)
from zipline.finance.trading import SimulationParameters
from zipline.finance import trading
from zipline.sources.test_source import (
    date_gen,
    create_trade
)

# For backwards compatibility
from zipline.data.loader import (load_from_yahoo,
                                 load_bars_from_yahoo)

__all__ = ['load_from_yahoo', 'load_bars_from_yahoo']


def create_simulation_parameters(year=2006, start=None, end=None,
                                 capital_base=float("1.0e5"),
                                 num_days=None, load=None
                                 ):
    """Construct a complete environment with reasonable defaults"""
    if start is None:
        start = datetime(year, 1, 1, tzinfo=pytz.utc)
    if end is None:
        if num_days:
            trading.environment = trading.TradingEnvironment(load=load)
            start_index = trading.environment.trading_days.searchsorted(
                start)
            end = trading.environment.trading_days[start_index + num_days - 1]
        else:
            end = datetime(year, 12, 31, tzinfo=pytz.utc)
    sim_params = SimulationParameters(
        period_start=start,
        period_end=end,
        capital_base=capital_base,
    )

    return sim_params


def create_noop_environment():
    oneday = timedelta(days=1)
    start = datetime(2006, 1, 1, tzinfo=pytz.utc)

    days = []
    tr_curves = OrderedDict()
    for day in date_gen(start=start, delta=oneday, count=252):
        days.append(day)
        curve = {
            '10year': 0.0799,
            '1month': 0.0799,
            '1year': 0.0785,
            '20year': 0.0765,
            '2year': 0.0794,
            '30year': 0.0804,
            '3month': 0.0789,
            '3year': 0.0796,
            '5year': 0.0792,
            '6month': 0.0794,
            '7year': 0.0804,
            'tid': 1752
        }
        tr_curves[day] = curve

    bm_returns = pd.Series(index=days, data=0.1)

    load_nodata = lambda x: (bm_returns, tr_curves)

    return trading.TradingEnvironment(load=load_nodata)


def create_random_simulation_parameters():
    trading.environment = trading.TradingEnvironment()
    treasury_curves = trading.environment.treasury_curves

    for n in range(100):

        random_index = random.randint(
            0,
            len(treasury_curves) - 1
        )

        start_dt = treasury_curves.index[random_index]
        end_dt = start_dt + timedelta(days=365)

        now = datetime.utcnow().replace(tzinfo=pytz.utc)

        if end_dt <= now:
            break

    assert end_dt <= now, """
failed to find a suitable daterange after 100 attempts. please double
check treasury and benchmark data in findb, and re-run the test."""

    sim_params = SimulationParameters(
        period_start=start_dt,
        period_end=end_dt
    )

    return sim_params, start_dt, end_dt


def get_next_trading_dt(current, interval):
    next_dt = pd.Timestamp(current).tz_convert(trading.environment.exchange_tz)

    while True:
        # Convert timestamp to naive before adding day, otherwise the when
        # stepping over EDT an hour is added.
        next_dt = pd.Timestamp(next_dt.replace(tzinfo=None))
        next_dt = next_dt + interval
        next_dt = pd.Timestamp(next_dt, tz=trading.environment.exchange_tz)
        next_dt_utc = next_dt.tz_convert('UTC')
        if trading.environment.is_market_hours(next_dt):
            break
        next_dt = next_dt_utc.tz_convert(trading.environment.exchange_tz)

    return next_dt_utc


def create_trade_history(sid, prices, amounts, interval, sim_params,
                         source_id="test_factory"):
    trades = []
    current = sim_params.first_open

    oneday = timedelta(days=1)
    use_midnight = interval >= oneday
    for price, amount in zip(prices, amounts):
        if use_midnight:
            trade_dt = current.replace(hour=0, minute=0)
        else:
            trade_dt = current
        trade = create_trade(sid, price, amount, trade_dt, source_id)
        trades.append(trade)
        current = get_next_trading_dt(current, interval)

    assert len(trades) == len(prices)
    return trades
{'fraction_non_alphanumeric': 0.08578695099285924, 'fraction_numerical': 0.023672111904529003, 'mean_word_length': 2.5218739235273855, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3916899', 'n_tokens_mistral': 3271, 'n_tokens_neox': 3004, 'n_words': 789}