Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Enhancement] Get Field Descriptions from OBBject #6117

Open
wants to merge 13 commits into
base: develop
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 11 additions & 1 deletion openbb_platform/core/integration/test_obbject.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,7 @@ def test_to_chart(obb):
"""Test obbject to chart."""

res = obb.equity.price.historical("AAPL", provider="fmp")
res.charting.to_chart()
res.charting.to_chart(render=False)
assert isinstance(res.chart.fig, OpenBBFigure)


Expand All @@ -89,3 +89,13 @@ def test_show(obb):
stocks_data = obb.equity.price.historical("AAPL", provider="fmp", chart=True)
assert isinstance(stocks_data.chart.fig, OpenBBFigure)
assert stocks_data.chart.fig.show() is None


@pytest.mark.integration
def test_get_field_descriptions(obb):
"""Test obbject get field descriptions."""

obb_data = obb.equity.profile("MSFT", provider="yfinance")
descriptions = obb_data.get_field_descriptions()
assert isinstance(descriptions, dict)
assert len(obb_data.to_df(index=None).columns) == len(descriptions)
26 changes: 26 additions & 0 deletions openbb_platform/core/openbb_core/app/model/obbject.py
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,32 @@ def model_parametrized_name(cls, params: Any) -> str:
"""Return the model name with the parameters."""
return f"OBBject[{cls.results_type_repr(params)}]"

def get_field_descriptions(self) -> Dict[str, str]:
"""
Get a dictionary of the returned field keys with their descriptions.

Fields automatically created by `alias_generator` will not have descriptions.
"""
descriptions = {}
model = None
if isinstance(self.results, list):
model = self.results[0].model_json_schema(by_alias=False).get("properties", None) # type: ignore
columns = self.to_df(index=None).columns.to_list() # type: ignore
if columns[0] == 0:
columns = self.to_df(index=None).iloc[:, 0].to_list()
else:
model = self.results.model_json_schema(by_alias=False).get("properties", None) # type: ignore
columns = list(self.results.model_dump(exclude_none=True).keys()) # type: ignore
if model is None:
raise OpenBBError(
"Could not extract model property definitions from OBBject."
)
for i in columns:
if i in model:
descriptions[i] = model[i].get("description", None)

return descriptions

def to_df(
self, index: Optional[Union[str, None]] = "date", sort_by: Optional[str] = None
) -> pd.DataFrame:
Expand Down
63 changes: 63 additions & 0 deletions openbb_platform/core/tests/app/model/test_obbject.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
"""Tests for the OBBject class."""

from typing import Optional
from unittest.mock import MagicMock

import pandas as pd
Expand All @@ -8,6 +9,7 @@
from openbb_core.app.utils import basemodel_to_df
from openbb_core.provider.abstract.data import Data
from pandas.testing import assert_frame_equal
from pydantic import Field


def test_OBBject():
Expand Down Expand Up @@ -56,6 +58,22 @@ class MockDataFrame(Data):
value: float


class MockDataModel(Data):

name: str = Field(default=None, description="Common name of the company.")
cik: str = Field(
default=None,
description="Central Index Key assigned to the company.",
)
cusip: str = Field(default=None, description="CUSIP identifier for the company.")
isin: str = Field(
default=None, description="International Securities Identification Number."
)
lei: Optional[str] = Field(
default=None, description="Legal Entity Identifier assigned to the company."
)


@pytest.mark.parametrize(
"results, expected_df",
[
Expand Down Expand Up @@ -388,3 +406,48 @@ def test_show_chart_no_fig():
# Act and Assert
with pytest.raises(OpenBBError, match="Chart not found."):
mock_instance.show()


@pytest.mark.parametrize(
"results",
[
# Test case 1: List of models.
(
[
MockDataModel(
name="Mock Company",
cik="0001234567",
cusip="5556789",
isin="US5556789",
lei=None,
),
MockDataModel(
name="Mock Company 2",
cik="0001234568",
cusip="5556781",
isin="US5556788",
lei="1234567890",
),
]
),
# Test case 2: Not a list.
MockDataModel(
name="Mock Company 3",
cik="0001234565",
cusip="5556783",
isin="US5556785",
lei="1234567891",
),
],
)
def test_get_field_descriptions(results):
"""Test helper."""
# Arrange
co = OBBject(results=results)

# Act
descriptions = co.get_field_descriptions()

# Assert
assert isinstance(descriptions, dict)
assert len(descriptions) == 5
Original file line number Diff line number Diff line change
Expand Up @@ -178,7 +178,7 @@ def to_chart(
index = (
data.index.name
if has_data and isinstance(data, (pd.DataFrame, pd.Series))
else ""
else None
)
data_as_df: pd.DataFrame = (
basemodel_to_df(convert_to_basemodel(data), index=index)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -181,7 +181,10 @@ def __plot__(

if not isinstance(indicators, ChartIndicators):
indicators = ChartIndicators.from_dict(indicators or dict(dict()))

df_stock.loc[:, "date"] = df_stock.index # type: ignore
df_stock["date"] = df_stock["date"].apply(pd.to_datetime)
df_stock.index = df_stock["date"] # type: ignore
df_stock.drop(columns=["date"], inplace=True)
self.indicators = indicators
self.intraday = df_stock.index[-2].time() != df_stock.index[-1].time()
self.df_stock = df_stock.sort_index(ascending=True)
Expand Down
Loading