A human interface for candle data in strategy.

jesse/strategies/Strategy.py

@@ -1,6 +1,7 @@
 from abc import ABC, abstractmethod
 from time import sleep
 from typing import List, Dict, Union
+from datetime import datetime
 
 import numpy as np
 
@@ -18,6 +19,70 @@
 from jesse.services.color import generate_unique_hex_color
 
 
+class StrategyCandles(np.ndarray):
+    """
+    Human view of candles in strategy.
+    """
+
+    def __new__(cls, input_array: np.ndarray, *args, **kwargs) -> "StrategyCandles":
+
+        if len(input_array.shape) != 2 or input_array.shape[1] != 6:
+            raise ValueError(f'Candle data size invalid: expected (*, 6), got {input_array.shape}')
+
+        return np.asarray(input_array, *args, **kwargs).view(cls)
+
+    def __getitem__(self, key: Union[int, slice]) -> Union[np.ndarray, "StrategyCandles", float]:
+
+        result = super().__getitem__(key)
+
+        if isinstance(result, np.ndarray):
+            if result.shape == (6,) or (len(result.shape) == 2 and result.shape[1] == 6):
+                return result # remains this class
+
+            return result.view(np.ndarray)
+
+        return result
+
+    @property
+    def time(self) -> Union[np.ndarray, float]:
+        return self[0] if self.shape == (6,) else self[:, 0]
+
+    @property
+    def open(self) -> Union[np.ndarray, float]:
+        return self[1] if self.shape == (6,) else self[:, 1]
+
+    @property
+    def close(self) -> Union[np.ndarray, float]:
+        return self[2] if self.shape == (6,) else self[:, 2]
+
+    @property
+    def high(self) -> Union[np.ndarray, float]:
+        return self[3] if self.shape == (6,) else self[:, 3]
+
+    @property
+    def low(self) -> Union[np.ndarray, float]:
+        return self[4] if self.shape == (6,) else self[:, 4]
+
+    @property
+    def volume(self) -> Union[np.ndarray, float]:
+        return self[5] if self.shape == (6,) else self[:, 5]
+
+    @property
+    def time_dt(self) -> Union[datetime, List[datetime]]:
+
+        if self.shape == (6,):
+            return datetime.fromtimestamp(self[0] / 1000)
+
+        return [datetime.fromtimestamp(time / 1000) for time in self[:, 0]]
+
+    def __repr__(self) -> str:
+
+        if self.shape == (6,):
+            return f'{self.time_dt} open {self.open} close {self.close} high {self.high} low {self.low}'
+
+        return f'{len(self)} candles from {self[0].time_dt} to {self[-1].time_dt}'
+
+
 class Strategy(ABC):
     """
     The parent strategy class which every strategy must extend. It is the heart of the framework!
@@ -1114,13 +1179,13 @@ def low(self) -> float:
         return self.current_candle[4]
 
     @property
-    def candles(self) -> np.ndarray:
+    def candles(self) -> StrategyCandles:
         """
         Returns candles for current trading route
 
         :return: np.ndarray
         """
-        return store.candles.get_candles(self.exchange, self.symbol, self.timeframe)
+        return StrategyCandles(store.candles.get_candles(self.exchange, self.symbol, self.timeframe))
 
     def get_candles(self, exchange: str, symbol: str, timeframe: str) -> np.ndarray:
         """

tests/test_strategy_candles.py

@@ -0,0 +1,55 @@
+import datetime as dt
+
+import numpy as np
+import pytest
+
+from jesse.strategies.Strategy import StrategyCandles
+
+def test_strategy_candles():
+
+    c1 = np.array([1543387200000, 200, 190, 220, 300, 1000])
+    c2 = np.array([1543387500000, 190, 250, 220, 301, 1010])
+    c3 = np.array([1543387800000, 250, 270, 220, 302, 1020])
+
+    candles_np = np.hstack((c1, c2, c3))
+    with pytest.raises(ValueError):
+        candles = StrategyCandles(candles_np)
+
+    candles_np = np.vstack((c1, c2, c3))
+    candles = StrategyCandles(candles_np)
+
+    assert candles.time[1] == 1543387500000
+    assert candles.open[1] == 190
+    assert candles.close[1] == 250
+    assert candles.high[1] == 220
+    assert candles.low[1] == 301
+    assert candles.time_dt[1] == dt.datetime(2018, 11, 28, 9, 45)
+
+    assert (candles.time == candles_np[:, 0]).all()
+    assert (candles.open == candles_np[:, 1]).all()
+    assert (candles.close == candles_np[:, 2]).all()
+    assert (candles.high == candles_np[:, 3]).all()
+    assert (candles.low == candles_np[:, 4]).all()
+    assert (candles.volume == candles_np[:, 5]).all()
+
+    assert (candles[0] == c1).all()
+    assert (candles[1] == c2).all()
+    assert (candles[2] == c3).all()
+
+    assert candles[-2].time == candles.time[-2]
+    assert candles[-2].open == candles.open[-2]
+    assert candles[-2].close == candles.close[-2]
+    assert candles[-2].high == candles.high[-2]
+    assert candles[-2].low == candles.low[-2]
+    assert candles[-2].volume == candles.volume[-2]
+    assert candles[-2].time_dt == candles.time_dt[-2]
+
+    assert candles.time_dt[-1] == dt.datetime.fromtimestamp(candles.time[-1] / 1000)
+
+    assert type(candles[0]) == type(candles)
+    assert type(candles[0:1]) == type(candles)
+    assert type(candles[0:1, 0:2]) == np.ndarray
+
+
+
+