def is_same_input_shape(input: Optional[DataEntry],
input_shape: Optional[DataEntry]) -> bool:
if not input or not input_shape:
return False
return all([
input.data_type == input_shape.data_type,
time_util.to_time(input.timestamp) == time_util.to_time(
input_shape.timestamp),
])
def add(self, data_entry: data_entry_pb2.DataEntry):
data_key = key.make_key(data_entry)
data_entries = self._data[data_key]
if time_util.to_time(data_entry.timestamp) in {
time_util.to_time(e.timestamp)
for e in data_entries
}:
raise errors.AlreadyExistsError(
f'The data entry to add already exists: {data_entry}')
self._data[data_key].append(data_entry)
self._data[data_key].sort(key=lambda e: e.timestamp.ToSeconds())
def make_lookup_key(
data_entry: data_entry_pb2.DataEntry) -> store.DataStore.LookupKey:
return store.DataStore.LookupKey(data_space=data_entry.data_space,
symbol=data_entry.symbol,
data_type=data_entry.data_type,
timestamp=time_util.to_time(
data_entry.timestamp))
def test_read_data_that_does_not_exist(self):
self._store.add(get_data_entry_one())
with self.assertRaises(errors.NotFoundError):
lookup_key = store.DataStore.LookupKey(
data_space=int(_DataEntry.STOCK_DATA),
symbol="SPY",
data_type=data_type_pb2.DataType.CLOSE_PRICE,
timestamp=time_util.to_time(
timestamp_pb2.Timestamp(seconds=654321)))
self._store.read(lookup_key)
def read(
self,
lookup_key: store.DataStore.LookupKey) -> data_entry_pb2.DataEntry:
data_entries = self._data[key.from_lookup_key(lookup_key)]
try:
return next(
e for e in data_entries
if time_util.to_time(e.timestamp) == lookup_key.timestamp)
except StopIteration:
raise errors.NotFoundError(
f'Cannot find data matching lookup key: {lookup_key}')
def lookup(
self, lookup_key: store.DataStore.LookupKey
) -> data_entry_pb2.DataEntries:
results = []
for data_key, data_entries in self._data.items():
if not key.key_matches(data_key, lookup_key):
continue
results.extend(
(e for e in data_entries if lookup_key.timestamp is None
or time_util.to_time(e.timestamp) == lookup_key.timestamp))
# Return a value, as extend makes copies.
data_entries = data_entry_pb2.DataEntries()
data_entries.entries.extend(results)
return data_entries
def calculate(self, inputs: calc.CalcInputs) -> DataEntry:
if not inputs:
raise errors.InvalidArgumentError('Inputs cannot be empty')
symbol = input_util.are_for_stock(inputs)
t = time_util.to_time(inputs[-1].timestamp)
try:
input_util.validate_inputs(inputs, self.recursive_inputs_shape(t))
ema = self._recur(inputs)
except errors.InvalidArgumentError:
input_util.validate_inputs(inputs, self.source_inputs_shape(t))
ema = self._initial(inputs)
except errors.InvalidArgumentError as e:
raise e
return DataEntry(
symbol=symbol,
data_space=data_entry_pb2.DataEntry.STOCK_DATA,
data_type=self._calc_type,
value=ema,
timestamp=time_util.from_time(t),
)
def _make_input_shape_key(data_entry: DataEntry) -> Tuple[int, time.Time]:
return (data_entry.data_type, time_util.to_time(data_entry.timestamp))