mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-05 06:12:09 +08:00
added algotime get/set to algorithm
This commit is contained in:
@@ -12,10 +12,14 @@
|
||||
# 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.
|
||||
from copy import copy
|
||||
|
||||
import pytz
|
||||
import pandas as pd
|
||||
import numpy as np
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
from zipline.sources import DataFrameSource
|
||||
from zipline.utils.factory import create_trading_environment
|
||||
from zipline.transforms.utils import StatefulTransform
|
||||
@@ -63,6 +67,7 @@ class TradingAlgorithm(object):
|
||||
self.order = None
|
||||
self.frame_count = 0
|
||||
self.portfolio = None
|
||||
self.datetime = None
|
||||
|
||||
self.registered_transforms = {}
|
||||
self.transforms = []
|
||||
@@ -221,6 +226,22 @@ class TradingAlgorithm(object):
|
||||
def set_logger(self, logger):
|
||||
self.logger = logger
|
||||
|
||||
def set_datetime(self, dt):
|
||||
assert isinstance(dt, datetime), \
|
||||
"Attempt to set algorithm's current time with non-datetime"
|
||||
assert dt.tzinfo == pytz.utc, \
|
||||
"Algorithm expects a utc datetime"
|
||||
self.datetime = dt
|
||||
|
||||
def get_datetime(self):
|
||||
"""
|
||||
Returns a copy of the datetime.
|
||||
"""
|
||||
date_copy = copy(self.datetime)
|
||||
assert date_copy.tzinfo == pytz.utc, \
|
||||
"Algorithm should have a utc datetime"
|
||||
return date_copy
|
||||
|
||||
def init(self, *args, **kwargs):
|
||||
"""Called from constructor."""
|
||||
pass
|
||||
|
||||
@@ -269,6 +269,7 @@ class AlgorithmSimulator(object):
|
||||
# Needs to be set so that we inject the proper date into algo
|
||||
# log/print lines.
|
||||
self.snapshot_dt = date
|
||||
self.algo.set_datetime(self.snapshot_dt)
|
||||
start_tic = datetime.now()
|
||||
self.algo.handle_data(self.universe)
|
||||
stop_tic = datetime.now()
|
||||
|
||||
Reference in New Issue
Block a user