Coverage for source/data_handling/data_handler.py: 95%

1# data_handling/data_handler.py 

2 

3# global imports 

4import pandas as pd 

5from typing import Optional 

6 

7# local imports 

8from source.data_handling import CoinBaseHandler 

9from source.indicators import IndicatorHandlerBase 

10from source.utils import Granularity 

11 

12class DataHandler(): 

13 """ 

14 Responsible for data handling. Including data collection and preparation. 

15 """ 

16 

17 def __init__(self, list_of_indicators_to_apply: Optional[list[IndicatorHandlerBase]] = None) -> None: 

18 """ 

19 Class constructor. 

20 

21 Parameters: 

22 list_of_indicators_to_apply (list): List of indicators further to apply. 

23 """ 

24 

25 if list_of_indicators_to_apply is None: 

26 list_of_indicators_to_apply = [] 

27 self.__indicators: list[IndicatorHandlerBase] = list_of_indicators_to_apply 

28 self.__coinbase: CoinBaseHandler = CoinBaseHandler() 

29 

30 async def prepare_data(self, trading_pair: str, start_date: str, end_date: str, granularity: Granularity) -> pd.DataFrame: 

31 """ 

32 Collects data from coinbase API and extends it with assigned list of indicators. 

33 

34 Parameters: 

35 trading_pair (str): String representing unique trainding pair symbol. 

36 start_date (str): String representing date that collected data should start from. 

37 end_date (str): String representing date that collected data should finish at. 

38 granularity (Granularity): Enum specifying resolution of collected data - e.g. each 

39 15 minutes or 1 hour or 6 hours is treated separately 

40 

41 Raises: 

42 RuntimeError: If given traiding pair symbol is not recognized. 

43 

44 Returns: 

45 (pd.DataFrame): Collected data extended with given indicators. 

46 """ 

47 

48 possible_traiding_pairs = await self.__coinbase.get_possible_pairs() 

49 if trading_pair not in possible_traiding_pairs.index: 

50 raise RuntimeError('Traiding pair not recognized!') 

51 

52 data = await self.__coinbase.get_candles_for(trading_pair, start_date, end_date, granularity) 

53 if len(self.__indicators) > 0: 

54 indicators_data = [] 

55 for indicator in self.__indicators: 

56 indicators_data.append(indicator.calculate(data)) 

57 data = pd.concat([data] + indicators_data, axis=1) 

58 

59 return data