Coverage for adrs/data/__init__.py: 71%

1from __future__ import annotations 

2 

3from typing import TYPE_CHECKING 

4from datetime import datetime, timedelta 

5 

6from .processor import DataProcessor 

7 

8from .types import DataInfo, DataColumn 

9from .datamap import Datamap 

10from .dataloader import DataLoader 

11from .handler import yfinance_handler 

12from .cache import Cache 

13from .connector import DatasourceStream, MetricStream, MetricBuilder 

14 

15if TYPE_CHECKING: 

16 from adrs.performance import Evaluator 

17 

18__all__ = [ 

19 "Datamap", 

20 "DataProcessor", 

21 "DataInfo", 

22 "DataColumn", 

23 "DataLoader", 

24 "yfinance_handler", 

25 "MetricStream", 

26 "MetricBuilder", 

27 "DatasourceStream", 

28 "Cache", 

29 "DataLoader", 

30] 

31 

32 

33async def make_datamap( 

34 dataloader: DataLoader, 

35 start_time: datetime, 

36 end_time: datetime, 

37 data_infos: list[DataInfo], 

38 evaluator: Evaluator | None = None, 

39 evaluator_offset: timedelta = timedelta(days=1), 

40) -> Datamap: 

41 """Create a datamap and initialize with given data infos and optionally evaluator.""" 

42 

43 datamap = Datamap() 

44 

45 # Setup the datamap (download data) 

46 await datamap.init( 

47 dataloader=dataloader, 

48 infos=data_infos, 

49 start_time=start_time, 

50 end_time=end_time, 

51 ) 

52 

53 # download data with (+1 day offset for candle shift) 

54 if evaluator is not None: 

55 await datamap.init( 

56 dataloader=dataloader, 

57 infos=list(evaluator.assets.values()), 

58 start_time=start_time, 

59 end_time=end_time + evaluator_offset, 

60 ) 

61 

62 return datamap