Coverage for adrs/data/handler.py: 38%

1import logging 

2import pandas as pd 

3import polars as pl 

4import yfinance as yf 

5 

6from datetime import datetime 

7 

8from adrs.types import Topic 

9from adrs.data.cache import Cache 

10from adrs.data.datasource import Datasource 

11 

12 

13logger = logging.getLogger(__name__) 

14 

15 

16def datasource_handler(datasource: Datasource, cache: Cache): 

17 async def handler(topic_str: str, start_time: datetime, end_time: datetime): 

18 topic = Topic.from_str(topic_str) 

19 

20 return await cache.fetch( 

21 datasource=datasource, 

22 topic=topic, 

23 start_time=start_time, 

24 end_time=end_time, 

25 ) 

26 

27 return handler 

28 

29 

30async def yfinance_handler(topic: str, start_time: datetime, end_time: datetime): 

31 _topic = Topic.from_str(topic) 

32 if _topic.provider != "yfinance": 

33 return 

34 

35 match _topic.endpoint: 

36 case "candle": 

37 ticker = _topic.query_params().get("ticker") 

38 if not ticker: 

39 raise ValueError(f"Topic {_topic} is missing 'ticker'") 

40 

41 interval = _topic.query_params().get("interval", "1d") 

42 

43 ticker = yf.Ticker(ticker) 

44 df: pd.DataFrame = ticker.history( 

45 start=start_time.strftime("%Y-%m-%d"), 

46 end=end_time.strftime("%Y-%m-%d"), 

47 interval=interval, 

48 ) 

49 return pl.from_pandas(df.reset_index()).select( 

50 pl.col("Date").alias("start_time").dt.convert_time_zone("UTC"), 

51 pl.col("Open").alias("open"), 

52 pl.col("High").alias("high"), 

53 pl.col("Low").alias("low"), 

54 pl.col("Close").alias("close"), 

55 pl.col("Volume").alias("volume").cast(pl.Float64), 

56 ) 

57 case _: 

58 raise ValueError(f"Unsupported endpoint: {_topic.endpoint}")