From 57deb8eeab1e7763119efdeb8e73127927601044 Mon Sep 17 00:00:00 2001 From: montezdesousa <79287829+montezdesousa@users.noreply.github.com> Date: Wed, 3 Apr 2024 21:46:09 +0100 Subject: [PATCH] [Feature] - Annotated results (#6282) * feat: pass metadata through AnnotatedData * fix: typing * fix: rename AnnotatedData to AnnotatedResult * fix: mypy --- .../core/openbb_core/app/model/obbject.py | 8 +++- .../provider/abstract/annotated_result.py | 20 +++++++++ .../openbb_core/provider/abstract/fetcher.py | 15 ++++--- .../fred/openbb_fred/models/series.py | 41 ++++++++++--------- 4 files changed, 58 insertions(+), 26 deletions(-) create mode 100644 openbb_platform/core/openbb_core/provider/abstract/annotated_result.py diff --git a/openbb_platform/core/openbb_core/app/model/obbject.py b/openbb_platform/core/openbb_core/app/model/obbject.py index 89232702405..8d4ba26839c 100644 --- a/openbb_platform/core/openbb_core/app/model/obbject.py +++ b/openbb_platform/core/openbb_core/app/model/obbject.py @@ -26,6 +26,7 @@ from openbb_core.app.model.abstract.tagged import Tagged from openbb_core.app.model.abstract.warning import Warning_ from openbb_core.app.model.charts.chart import Chart from openbb_core.app.utils import basemodel_to_df +from openbb_core.provider.abstract.annotated_result import AnnotatedResult from openbb_core.provider.abstract.data import Data if TYPE_CHECKING: @@ -299,4 +300,9 @@ class OBBject(Tagged, Generic[T]): OBBject[ResultsType] OBBject with results. """ - return cls(results=await query.execute()) + results = await query.execute() + if isinstance(results, AnnotatedResult): + return cls( + results=results.result, extra={"results_metadata": results.metadata} + ) + return cls(results=results) diff --git a/openbb_platform/core/openbb_core/provider/abstract/annotated_result.py b/openbb_platform/core/openbb_core/provider/abstract/annotated_result.py new file mode 100644 index 00000000000..8d51086017b --- /dev/null +++ b/openbb_platform/core/openbb_core/provider/abstract/annotated_result.py @@ -0,0 +1,20 @@ +"""Annotated result.""" + +from typing import Generic, Optional, TypeVar + +from pydantic import BaseModel, Field + +T = TypeVar("T") + + +class AnnotatedResult(BaseModel, Generic[T]): + """Annotated result allows fetchers to return metadata along with the data.""" + + result: Optional[T] = Field( + default=None, + description="Serializable results.", + ) + metadata: Optional[dict] = Field( + default=None, + description="Metadata.", + ) diff --git a/openbb_platform/core/openbb_core/provider/abstract/fetcher.py b/openbb_platform/core/openbb_core/provider/abstract/fetcher.py index 011f03f197b..f2cc759821a 100644 --- a/openbb_platform/core/openbb_core/provider/abstract/fetcher.py +++ b/openbb_platform/core/openbb_core/provider/abstract/fetcher.py @@ -9,12 +9,14 @@ from typing import ( Generic, Optional, TypeVar, + Union, get_args, get_origin, ) from pandas import DataFrame +from openbb_core.provider.abstract.annotated_result import AnnotatedResult from openbb_core.provider.abstract.data import Data from openbb_core.provider.abstract.query_params import QueryParams from openbb_core.provider.utils.helpers import maybe_coroutine, run_async @@ -56,7 +58,7 @@ class Fetcher(Generic[Q, R]): """Extract the data from the provider.""" @staticmethod - def transform_data(query: Q, data: Any, **kwargs) -> R: + def transform_data(query: Q, data: Any, **kwargs) -> Union[R, AnnotatedResult[R]]: """Transform the provider-specific data.""" raise NotImplementedError @@ -79,7 +81,7 @@ class Fetcher(Generic[Q, R]): params: Dict[str, Any], credentials: Optional[Dict[str, str]] = None, **kwargs, - ) -> R: + ) -> Union[R, AnnotatedResult[R]]: """Fetch data from a provider.""" query = cls.transform_query(params=params) data = await maybe_coroutine( @@ -139,7 +141,7 @@ class Fetcher(Generic[Q, R]): data = run_async( cls.extract_data, query=query, credentials=credentials, **kwargs ) - transformed_data = cls.transform_data(query=query, data=data, **kwargs) + result = cls.transform_data(query=query, data=data, **kwargs) # Class Assertions assert isinstance( @@ -184,10 +186,13 @@ class Fetcher(Generic[Q, R]): assert len(data) > 0, "Data must not be empty." # Transformed Data Assertions + transformed_data = ( + result.result if isinstance(result, AnnotatedResult) else result + ) + assert transformed_data, "Transformed data must not be None." - is_list = isinstance(transformed_data, list) - if is_list: + if isinstance(transformed_data, list): return_type_args = cls.return_type.__args__[0] return_type_is_dict = ( hasattr(return_type_args, "__origin__") diff --git a/openbb_platform/providers/fred/openbb_fred/models/series.py b/openbb_platform/providers/fred/openbb_fred/models/series.py index 35db3d0a2a5..7c1108e73f3 100644 --- a/openbb_platform/providers/fred/openbb_fred/models/series.py +++ b/openbb_platform/providers/fred/openbb_fred/models/series.py @@ -1,10 +1,9 @@ """FRED Series Model.""" -import json -import warnings from typing import Any, Dict, List, Literal, Optional import pandas as pd +from openbb_core.provider.abstract.annotated_result import AnnotatedResult from openbb_core.provider.abstract.fetcher import Fetcher from openbb_core.provider.standard_models.fred_series import ( SeriesData, @@ -19,8 +18,6 @@ from openbb_core.provider.utils.helpers import ( ) from pydantic import Field -_warn = warnings.warn - class FredSeriesQueryParams(SeriesQueryParams): """FRED Series Query Params.""" @@ -145,9 +142,18 @@ class FredSeriesFetcher( f"{metadata_url}?series_id={series_id}&file_type=json&api_key={api_key}", timeout=5, ) - _metadata = metadata_response.get("seriess", [{}])[0] - observations = observations_response.get("observations") + # seriess is not a typo, it's the actual key in the response + _metadata = ( + metadata_response.get("seriess", [{}])[0] + if isinstance(metadata_response, dict) + else {} + ) or {} + observations = ( + observations_response.get("observations") + if isinstance(observations_response, dict) + else [] + ) or [] try: for d in observations: d.pop("realtime_start") @@ -177,24 +183,18 @@ class FredSeriesFetcher( results = await amake_requests(urls, callback, timeout=5, **kwargs) - metadata, data = {}, {} - for item in results: - for series_id, result in item.items(): - data[series_id] = result.pop("data") - metadata[series_id] = result - - _warn(json.dumps(metadata)) - - return data + return results # pylint: disable=unused-argument @staticmethod def transform_data( - query: FredSeriesQueryParams, data: Dict, **kwargs: Any - ) -> List[FredSeriesData]: + query: FredSeriesQueryParams, data: List[Dict[str, Any]], **kwargs: Any + ) -> AnnotatedResult[List[FredSeriesData]]: """Transform data.""" - results = ( - pd.DataFrame(data) + series = {_id: s.pop("data", {}) for d in data for _id, s in d.items()} + metadata = {_id: m for d in data for _id, m in d.items()} + records = ( + pd.DataFrame(series) .filter(items=query.symbol.split(","), axis=1) .reset_index() .rename(columns={"index": "date"}) @@ -202,4 +202,5 @@ class FredSeriesFetcher( .replace("N/A", None) .to_dict("records") ) - return [FredSeriesData.model_validate(d) for d in results] + validated = [FredSeriesData.model_validate(r) for r in records] + return AnnotatedResult(result=validated, metadata=metadata)