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