Coverage for adrs/data/dataloader.py: 37%

1from pathlib import Path 

2 

3from polars import DataFrame 

4from datetime import datetime 

5from typing import Optional, Callable, Awaitable, Any 

6 

7from adrs.data.cache import Cache 

8from adrs.data.handler import datasource_handler 

9from adrs.data.datasource import ( 

10 DEFAULT_API_URL, 

11 Datasource, 

12 CybotradeDatasource, 

13 ClickhouseDatasource, 

14) 

15 

16type Handler = Callable[[str, datetime, datetime], Awaitable[DataFrame | None]] 

17 

18 

19def parse_credentials( 

20 credentials: dict[str, Any], cybotrade_api_url: str | None = None 

21) -> Datasource | None: 

22 match credentials.get("type"): 

23 case "clickhouse": 

24 return ClickhouseDatasource( 

25 **{k: v for k, v in credentials.items() if k != "type"} 

26 ) 

27 case None: 

28 if ( 

29 key := credentials.get("cybotrade_api_key") 

30 or credentials.get("datasource_api_key") 

31 ) is not None: 

32 return CybotradeDatasource( 

33 api_key=key, 

34 base_url=cybotrade_api_url 

35 if cybotrade_api_url 

36 else DEFAULT_API_URL, 

37 ) 

38 

39 

40class DataLoader: 

41 def __init__( 

42 self, 

43 data_dir: str, 

44 credentials: dict[str, Any], 

45 format: Optional[str] = None, 

46 cybotrade_api_url: Optional[str] = None, 

47 handlers: list[Handler] = [], 

48 ): 

49 self.data_dir = data_dir 

50 self.credentials = credentials 

51 self.format = format 

52 self.cybotrade_api_url = cybotrade_api_url 

53 self.handlers = handlers 

54 

55 datasource = parse_credentials(credentials, cybotrade_api_url) 

56 if datasource is None: 

57 raise Exception("invalid credentials format") 

58 self.datasource = datasource 

59 

60 async def load( 

61 self, 

62 topic: str, 

63 start_time: datetime, 

64 end_time: datetime, 

65 override_existing: bool = False, 

66 ) -> DataFrame: 

67 for handler in self.handlers: 

68 df = await handler(topic, start_time, end_time) 

69 if df is not None: 

70 return df 

71 

72 cache = Cache( 

73 data_path=Path(self.data_dir), 

74 format=self.format if self.format else "parquet", 

75 override_existing=override_existing, 

76 ) 

77 

78 default_handler = datasource_handler(datasource=self.datasource, cache=cache) 

79 return await default_handler( 

80 topic_str=topic, 

81 start_time=start_time, 

82 end_time=end_time, 

83 ) 

84 

85 def add_handler(self, handler: Handler): 

86 self.handlers.append(handler) 

87 

88 def remove_handler(self, handler: Handler): 

89 self.handlers = list(filter(lambda h: id(h) != id(handler), self.handlers))