Add filtering to AV historical price (#5928)

* add filtering to av historical price

* lints

---------

Co-authored-by: Pratyush Shukla <ps4534@nyu.edu>
This commit is contained in:
Henrique Joaquim
2024-01-09 15:05:58 +00:00
committed by GitHub
parent 16bd470a78
commit cf9536d869
2 changed files with 24 additions and 3 deletions

View File

@@ -5,7 +5,12 @@ from datetime import datetime
from typing import Any, Dict, List, Literal, Optional
from dateutil.relativedelta import relativedelta
from openbb_alpha_vantage.utils.helpers import extract_key_name, get_data, get_interval
from openbb_alpha_vantage.utils.helpers import (
extract_key_name,
filter_by_dates,
get_data,
get_interval,
)
from openbb_core.provider.abstract.fetcher import Fetcher
from openbb_core.provider.standard_models.equity_historical import (
EquityHistoricalData,
@@ -171,6 +176,7 @@ class AVEquityHistoricalFetcher(
return data[dynamic_key]
# pylint: disable=unused-argument
@staticmethod
def transform_data(
query: AVEquityHistoricalQueryParams, data: Dict, **kwargs: Any
@@ -180,5 +186,6 @@ class AVEquityHistoricalFetcher(
{"date": date, **{extract_key_name(k): v for k, v in values.items()}}
for date, values in data.items()
]
data = filter_by_dates(data, query.start_date, query.end_date)
return [AVEquityHistoricalData.model_validate(d) for d in data]

View File

@@ -2,8 +2,9 @@
import json
import re
from datetime import datetime
from io import StringIO
from typing import Any, List, Optional, TypeVar, Union
from typing import Any, Dict, List, Optional, TypeVar, Union
import requests
from openbb_core.provider import helpers
@@ -149,7 +150,6 @@ def get_data_one(url: str, **kwargs: Any) -> dict:
def get_interval(value: str) -> str:
"""Get the intervals for the Alpha Vantage API."""
intervals = {
"m": "min",
"d": "day",
@@ -164,3 +164,17 @@ def extract_key_name(key):
"""Extract the alphabetical part of the key using regex."""
match = re.search(r"\d+\.\s+([a-z]+)", key, re.I)
return match.group(1) if match else key
def filter_by_dates(
data: List[Dict[str, Any]], start_date: datetime, end_date: datetime
) -> List[Dict[str, Any]]:
"""Filter the data by start and end dates."""
return list(
filter(
lambda x: start_date
<= datetime.strptime(x["date"], "%Y-%m-%d").date()
<= end_date,
data,
)
)