def init_db(self, txn=None): """Connect to database and create tables. Parameters ---------- txn : sa.engine.Connection, optional The transaction to execute in. If this is not provided, a new transaction will be started with the engine provided. Returns ------- metadata : sa.MetaData The metadata that describes the new assets db. """ with ExitStack() as stack: if txn is None: txn = stack.enter_context(self.engine.begin()) tables_already_exist = self._all_tables_present(txn) # Create the SQL tables if they do not already exist. metadata.create_all(txn, checkfirst=True) if tables_already_exist: check_version_info(txn, version_info, ASSET_DB_VERSION) else: write_version_info(txn, version_info, ASSET_DB_VERSION)
def transform(self): """ Main generator work loop. """ algo = self.algo metrics_tracker = algo.metrics_tracker emission_rate = metrics_tracker.emission_rate def every_bar( dt_to_use, current_data=self.current_data, handle_data=algo.event_manager.handle_data, ): for capital_change in calculate_minute_capital_changes(dt_to_use): yield capital_change self.simulation_dt = dt_to_use # called every tick (minute or day). algo.on_dt_changed(dt_to_use) blotter = algo.blotter # handle any transactions and commissions coming out new orders # placed in the last bar new_transactions, new_commissions, closed_orders = blotter.get_transactions( current_data) blotter.prune_orders(closed_orders) for transaction in new_transactions: metrics_tracker.process_transaction(transaction) # since this order was modified, record it order = blotter.orders[transaction.order_id] metrics_tracker.process_order(order) for commission in new_commissions: metrics_tracker.process_commission(commission) handle_data(algo, current_data, dt_to_use) # grab any new orders from the blotter, then clear the list. # this includes cancelled orders. new_orders = blotter.new_orders blotter.new_orders = [] # if we have any new orders, record them so that we know # in what perf period they were placed. for new_order in new_orders: metrics_tracker.process_order(new_order) def once_a_day(midnight_dt, current_data=self.current_data, data_portal=self.data_portal): # process any capital changes that came overnight for capital_change in algo.calculate_capital_changes( midnight_dt, emission_rate=emission_rate, is_interday=True): yield capital_change # set all the timestamps self.simulation_dt = midnight_dt algo.on_dt_changed(midnight_dt) metrics_tracker.handle_market_open( midnight_dt, algo.data_portal, ) # handle any splits that impact any positions or any open orders. assets_we_care_about = (metrics_tracker.positions.keys() | algo.blotter.open_orders.keys()) if assets_we_care_about: splits = data_portal.get_splits(assets_we_care_about, midnight_dt) if splits: algo.blotter.process_splits(splits) metrics_tracker.handle_splits(splits) def on_exit(): # Remove references to algo, data portal, et al to break cycles # and ensure deterministic cleanup of these objects when the # simulation finishes. self.algo = None self.benchmark_source = self.current_data = self.data_portal = None with ExitStack() as stack: stack.callback(on_exit) stack.enter_context(self.processor) stack.enter_context(ZiplineAPI(self.algo)) if algo.data_frequency == "minute": def execute_order_cancellation_policy(): algo.blotter.execute_cancel_policy(SESSION_END) def calculate_minute_capital_changes(dt): # process any capital changes that came between the last # and current minutes return algo.calculate_capital_changes( dt, emission_rate=emission_rate, is_interday=False) else: def execute_order_cancellation_policy(): pass def calculate_minute_capital_changes(dt): return [] for dt, action in self.clock: if action == BAR: for capital_change_packet in every_bar(dt): yield capital_change_packet elif action == SESSION_START: for capital_change_packet in once_a_day(dt): yield capital_change_packet elif action == SESSION_END: # End of the session. positions = metrics_tracker.positions position_assets = algo.asset_finder.retrieve_all(positions) self._cleanup_expired_assets(dt, position_assets) execute_order_cancellation_policy() algo.validate_account_controls() yield self._get_daily_message(dt, algo, metrics_tracker) elif action == BEFORE_TRADING_START_BAR: self.simulation_dt = dt algo.on_dt_changed(dt) algo.before_trading_start(self.current_data) elif action == MINUTE_END: minute_msg = self._get_minute_message( dt, algo, metrics_tracker, ) yield minute_msg risk_message = metrics_tracker.handle_simulation_end( self.data_portal, ) yield risk_message
def ingest(name, environ=os.environ, timestamp=None, assets_versions=(), show_progress=False): """Ingest data for a given bundle. Parameters ---------- name : str The name of the bundle. environ : mapping, optional The environment variables. By default this is os.environ. timestamp : datetime, optional The timestamp to use for the load. By default this is the current time. assets_versions : Iterable[int], optional Versions of the assets db to which to downgrade. show_progress : bool, optional Tell the ingest function to display the progress where possible. """ try: bundle = bundles[name] except KeyError: raise UnknownBundle(name) calendar = get_calendar(bundle.calendar_name) start_session = bundle.start_session end_session = bundle.end_session if start_session is None or start_session < calendar.first_session: start_session = calendar.first_session if end_session is None or end_session > calendar.last_session: end_session = calendar.last_session if timestamp is None: timestamp = pd.Timestamp.utcnow() timestamp = timestamp.tz_convert('utc').tz_localize(None) timestr = to_bundle_ingest_dirname(timestamp) cachepath = cache_path(name, environ=environ) pth.ensure_directory(pth.data_path([name, timestr], environ=environ)) pth.ensure_directory(cachepath) with dataframe_cache(cachepath, clean_on_failure=False) as cache, \ ExitStack() as stack: # we use `cleanup_on_failure=False` so that we don't purge the # cache directory if the load fails in the middle if bundle.create_writers: wd = stack.enter_context( working_dir(pth.data_path([], environ=environ))) daily_bars_path = wd.ensure_dir( *daily_equity_relative(name, timestr)) daily_bar_writer = BcolzDailyBarWriter( daily_bars_path, calendar, start_session, end_session, ) # Do an empty write to ensure that the daily ctables exist # when we create the SQLiteAdjustmentWriter below. The # SQLiteAdjustmentWriter needs to open the daily ctables so # that it can compute the adjustment ratios for the dividends. daily_bar_writer.write(()) minute_bar_writer = BcolzMinuteBarWriter( wd.ensure_dir(*minute_equity_relative(name, timestr)), calendar, start_session, end_session, minutes_per_day=bundle.minutes_per_day, ) assets_db_path = wd.getpath(*asset_db_relative(name, timestr)) asset_db_writer = AssetDBWriter(assets_db_path) adjustment_db_writer = stack.enter_context( SQLiteAdjustmentWriter( wd.getpath(*adjustment_db_relative(name, timestr)), BcolzDailyBarReader(daily_bars_path), overwrite=True, )) else: daily_bar_writer = None minute_bar_writer = None asset_db_writer = None adjustment_db_writer = None if assets_versions: raise ValueError('Need to ingest a bundle that creates ' 'writers in order to downgrade the assets' ' db.') log.info("Ingesting {}.", name) bundle.ingest( environ, asset_db_writer, minute_bar_writer, daily_bar_writer, adjustment_db_writer, calendar, start_session, end_session, cache, show_progress, pth.data_path([name, timestr], environ=environ), ) for version in sorted(set(assets_versions), reverse=True): version_path = wd.getpath(*asset_db_relative( name, timestr, db_version=version, )) with working_file(version_path) as wf: shutil.copy2(assets_db_path, wf.path) downgrade(wf.path, version)
def ctx(self, *args, **kwargs): with ExitStack() as stack: for hook in self._hooks: sub_ctx = getattr(hook, method_name)(*args, **kwargs) stack.enter_context(sub_ctx) yield stack