-
Notifications
You must be signed in to change notification settings - Fork 844
/
data_handler.py
80 lines (62 loc) · 2.91 KB
/
data_handler.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
"""
Copyright 2019 Goldman Sachs.
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.
"""
import datetime as dt
from typing import Union
from gs_quant.backtests.data_sources import DataManager
from pytz import timezone, utc
class Clock(object):
def __init__(self):
self._time = None
self.reset()
def update(self, time: dt.datetime):
compare_time = self._time.replace(tzinfo=None) \
if time.tzinfo is None or time.tzinfo.utcoffset(time) is None else self._time
if time < compare_time:
raise RuntimeError(f'current time is {compare_time}, cannot run backwards to {time}')
self._time = time
def reset(self):
self._time = dt.datetime(1900, 1, 1).replace(tzinfo=utc)
def time_check(self, state: Union[dt.date, dt.datetime]):
if isinstance(state, dt.datetime):
if state.tzinfo is None or state.tzinfo.utcoffset(state) is None: # timezone naive
lookahead = state > self._time.replace(tzinfo=None)
else:
lookahead = state > self._time
else:
lookahead = state > self._time.date()
if lookahead:
raise RuntimeError(f'accessing data at {state} not allowed, current time is {self._time}')
class DataHandler(object):
def __init__(self, data_mgr: DataManager, tz: timezone):
self._data_mgr = data_mgr
self._clock = Clock()
self._tz = tz
def reset_clock(self):
self._clock.reset()
def update(self, state: dt.datetime):
self._clock.update(state)
def _utc_time(self, state: Union[dt.date, dt.datetime]):
# only switch to utc time if the datetime you've been sent is timezone naive
if isinstance(state, dt.datetime) and (state.tzinfo is None or state.tzinfo.utcoffset(state) is None):
return self._tz.localize(state).astimezone(utc).replace(tzinfo=None)
else:
return state
def get_data(self, state: Union[dt.date, dt.datetime], *key):
self._clock.time_check(state)
return self._data_mgr.get_data(self._utc_time(state), *key)
def get_data_range(self, start: Union[dt.date, dt.datetime], end: Union[dt.date, dt.datetime], *key):
self._clock.time_check(end)
if type(start) is not type(end):
raise RuntimeError('expect same type for start and end when asking for data range')
return self._data_mgr.get_data_range(self._utc_time(start), self._utc_time(end), *key)