def test_ago(self):
ts = DataSeriesTest.factory.build_time_series(series_id="test")
series = DataSeries(time_series=ts)
self.assertEqual(0.0, series.now())
series.add(data={"timestamp": 0, "v1": 1, "v2": 2})
self.assertEqual({"timestamp": 0, "v1": 1, "v2": 2}, series.ago(0))
series.add(data={"timestamp": 1, "v1": 1.2, "v2": 2.2})
self.assertEqual({"timestamp": 1, "v1": 1.2, "v2": 2.2}, series.ago(0))
self.assertEqual({"timestamp": 0, "v1": 1, "v2": 2}, series.ago(1))
series.add(data={"timestamp": 2, "v1": 1.3, "v2": 2.3})
self.assertEqual({"timestamp": 2, "v1": 1.3, "v2": 2.3}, series.ago(0))
self.assertEqual({"timestamp": 1, "v1": 1.2, "v2": 2.2}, series.ago(1))
self.assertEqual({"timestamp": 0, "v1": 1, "v2": 2}, series.ago(2))
series.add(data={"timestamp": 3, "v1": 1.4, "v2": 2.4})
self.assertEqual({"timestamp": 3, "v1": 1.4, "v2": 2.4}, series.ago(0))
self.assertEqual({"timestamp": 2, "v1": 1.3, "v2": 2.3}, series.ago(1))
self.assertEqual({"timestamp": 1, "v1": 1.2, "v2": 2.2}, series.ago(2))
self.assertEqual({"timestamp": 0, "v1": 1, "v2": 2}, series.ago(3))
self.assertEqual({"v1": 1.4, "v2": 2.4}, series.ago(0, ["v1", "v2"]))
self.assertEqual(1.4, series.ago(0, "v1"))
self.assertEqual(1.4, series.ago(0, ["v1"]))
class PerformanceAnalyzer(Analyzer):
Performance = "Performance"
StockValue = "stock_value"
Cash = "cash"
TotalEquity = "total_equity"
def __init__(self, portfolio, state):
self.portfolio = portfolio
self.state = state
self.series = DataSeries(time_series=self.state.performance.series)
def update(self, timestamp: int, total_equity: float):
self.state.performance.total_equity = total_equity
self.series.add(timestamp=timestamp,
data={
self.StockValue: self.state.stock_value,
self.Cash: self.state.cash,
self.TotalEquity: total_equity
})
def get_result(self):
return {
self.StockValue: self.state.stock_value,
self.Cash: self.state.cash,
self.TotalEquity: self.state.performance.total_equity
}
def get_series(self, keys=None):
keys = keys if keys else [self.StockValue, self.Cash, self.TotalEquity]
return self.series.get_series(keys)
def now(self, key):
return self.series.now(key)
def total_equity(self) -> float:
return self.state.performance.total_equity