jerry/OpenBB
1
0
Fork 0
mirror of https://github.com/OpenBB-finance/OpenBB.git synced 2026-11-01 15:01:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #5430 from OpenBB-finance/feature/improved_helpers

Browse Source 
Feature/improved helpers
This commit is contained in:
Igor Radovanovic
2023-09-12 15:25:38 +02:00
committed by GitHub
parent ceddb03fe6 d5ad4652c6
commit d7da32483a
3 changed files with 129 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
openbb_sdk/sdk/core/openbb_core/app/model/obbject.py
View File

@@ -97,6 +97,8 @@ class OBBject(GenericModel, Generic[T], Tagged):
res = self.results
df = None
sort_columns = True
if isinstance(res, list) and len(res) == 1 and isinstance(res[0], dict):
for r in res:
dict_of_df = {}
@@ -105,21 +107,25 @@ class OBBject(GenericModel, Generic[T], Tagged):
isinstance(nv, BaseModel) for nv in v
):
dict_of_df[k] = basemodel_to_df(v, "date")
sort_columns = False
else:
dict_of_df[k] = pd.DataFrame(v)
df = pd.concat(dict_of_df, axis=1)
elif isinstance(res, list) and all(
isinstance(item, BaseModel) for item in res
):
df = basemodel_to_df(res, "date")
sort_columns = False
else:
df = pd.DataFrame(res)
if df is None:
raise OpenBBError("Unsupported data format.")
# Improve output so that all columns that are None are not returned
df.sort_index(axis=1, inplace=True)
# Drop columns that are all NaN, but don't rearrange columns
if sort_columns:
df.sort_index(axis=1, inplace=True)
df = df.dropna(axis=1, how="all")
except Exception as e:
@@ -150,11 +156,27 @@ class OBBject(GenericModel, Generic[T], Tagged):
Dict[str, List]
Dictionary of lists.
"""
if isinstance(self.results, dict) and all(
isinstance(v, dict) for v in self.results.values()
):
results: Dict[str, List] = {}
for _, v in self.results.items():
for kk, vv in v.items():
if kk not in results:
results[kk] = []
results[kk].append(vv)
return results
df = self.to_dataframe().reset_index() # type: ignore
results = {}
for field in df.columns:
results[field] = df[field].tolist()
# remove index from results
if "index" in results:
del results["index"]
return results
def to_chart(self, **kwargs):

2
openbb_sdk/sdk/core/openbb_core/app/utils.py
View File

@@ -10,7 +10,6 @@ def basemodel_to_df(
index: Optional[Union[str, Iterable]] = None,
) -> pd.DataFrame:
"""Convert list of BaseModel to a Pandas DataFrame."""
if isinstance(data, list):
df = pd.DataFrame([d.dict() for d in data])
else:
@@ -30,6 +29,7 @@ def basemodel_to_df(
if df.index.name == "date":
df.index = pd.to_datetime(df.index)
df.sort_index(axis=0, inplace=True)
return df

158
openbb_sdk/sdk/core/tests/app/model/test_obbject.py
View File

@@ -5,6 +5,7 @@ import pandas as pd
import pytest
from openbb_core.app.model.obbject import Chart, OBBject, OpenBBError
from openbb_provider.abstract.data import Data
from pandas.testing import assert_frame_equal
def test_OBBject():
@@ -30,6 +31,21 @@ def test_to_dataframe_no_results():
co.to_dataframe()
class MockData(Data):
"""Test helper."""
x: int
y: int
class MockMultiData(Data):
"""Test helper."""
date: str
another_date: str
value: float
@pytest.mark.parametrize(
"results, expected_df",
[
@@ -51,24 +67,28 @@ def test_to_dataframe_no_results():
# Test case 3: List of Data
(
[
Data(x=0, y=2),
Data(x=1, y=3),
Data(x=2, y=0),
Data(x=3, y=1),
Data(x=4, y=6),
MockData(x=0, y=2),
MockData(x=1, y=3),
MockData(x=2, y=0),
MockData(x=3, y=1),
MockData(x=4, y=6),
],
pd.DataFrame({"x": [0, 1, 2, 3, 4], "y": [2, 3, 0, 1, 6]}),
pd.DataFrame(
{"x": [0, 1, 2, 3, 4], "y": [2, 3, 0, 1, 6]}, columns=["x", "y"]
),
),
# Test case 4: List of dict
(
[
{"x": 1, "y": 2},
{"x": 1, "y": 3},
{"x": 2, "y": 0},
{"x": 3, "y": 1},
{"x": 4, "y": 6},
{"a": 1, "y": 2},
{"a": 1, "y": 3},
{"a": 2, "y": 0},
{"a": 3, "y": 1},
{"a": 4, "y": 6},
],
pd.DataFrame({"x": [1, 1, 2, 3, 4], "y": [2, 3, 0, 1, 6]}),
pd.DataFrame(
{"a": [1, 1, 2, 3, 4], "y": [2, 3, 0, 1, 6]}, columns=["a", "y"]
),
),
# Test case 5: List of Lists
(
@@ -121,14 +141,26 @@ def test_to_dataframe_no_results():
[
{
"df1": [
Data(date="1956-01-01", another_date="2023-09-01", value=0.0),
Data(date="1956-02-01", another_date="2023-09-01", value=0.0),
Data(date="1956-03-01", another_date="2023-09-01", value=0.0),
MockMultiData(
date="1956-01-01", another_date="2023-09-01", value=0.0
),
MockMultiData(
date="1956-02-01", another_date="2023-09-01", value=0.0
),
MockMultiData(
date="1956-03-01", another_date="2023-09-01", value=0.0
),
],
"df2": [
Data(date="1955-03-01", another_date="2023-09-01", value=0.0),
Data(date="1955-04-01", another_date="2023-09-01", value=0.0),
Data(date="1955-05-01", another_date="2023-09-01", value=0.0),
MockMultiData(
date="1955-03-01", another_date="2023-09-01", value=0.0
),
MockMultiData(
date="1955-04-01", another_date="2023-09-01", value=0.0
),
MockMultiData(
date="1955-05-01", another_date="2023-09-01", value=0.0
),
],
}
],
@@ -141,13 +173,10 @@ def test_to_dataframe_no_results():
pd.to_datetime("1956-02-01"),
pd.to_datetime("1956-03-01"),
],
"another_date": [
"2023-09-01",
"2023-09-01",
"2023-09-01",
],
"another_date": ["2023-09-01", "2023-09-01", "2023-09-01"],
"value": [0.0, 0.0, 0.0],
}
},
columns=["date", "another_date", "value"],
).set_index("date"),
"df2": pd.DataFrame(
{
@@ -156,16 +185,14 @@ def test_to_dataframe_no_results():
pd.to_datetime("1955-04-01"),
pd.to_datetime("1955-05-01"),
],
"another_date": [
"2023-09-01",
"2023-09-01",
"2023-09-01",
],
"another_date": ["2023-09-01", "2023-09-01", "2023-09-01"],
"value": [0.0, 0.0, 0.0],
}
},
columns=["date", "another_date", "value"],
).set_index("date"),
},
axis=1,
sort=True,
),
),
# Test case 10: Empty results
@@ -182,7 +209,7 @@ def test_to_dataframe(results, expected_df):
# Act and Assert
if isinstance(expected_df, pd.DataFrame):
result = co.to_dataframe()
assert result.equals(expected_df)
assert_frame_equal(result, expected_df)
else:
with pytest.raises(expected_df.__class__) as exc_info:
co.to_dataframe()
@@ -191,45 +218,68 @@ def test_to_dataframe(results, expected_df):
@pytest.mark.parametrize(
"results, expected",
[
# Test case 1: Normal results with "date" column
"results, expected_dict",
[ # Case 1: Normal results with "date" column
(
[{"date": "2023-07-30", "value": 10}, {"date": "2023-07-31", "value": 20}],
{
"date": [pd.to_datetime("2023-07-30"), pd.to_datetime("2023-07-31")],
"value": [10, 20],
},
{"date": ["2023-07-30", "2023-07-31"], "value": [10, 20]},
),
# Test case 2: Normal results without "date" column
# Case 2: Normal results without "date" column
(
[{"value": 10}, {"value": 20}],
{
"value": [10, 20],
},
{"value": [10, 20]},
),
# Test case 3: Empty results
# Test case 3: Dict of lists
(
{"0": [0, 2], "1": [1, 3], "2": [2, 0], "3": [3, 1], "4": [4, 6]},
{"0": [0, 2], "1": [1, 3], "2": [2, 0], "3": [3, 1], "4": [4, 6]},
),
# Test case 4: No results
([], OpenBBError("Results not found.")),
# Test case 5: Results as None, should raise OpenBBError
(None, OpenBBError("Results not found.")),
# Test case 6: List of tuples
(
[(3, 2), (1, 3), (2, 0), (3, 1), (4, 6)],
{0: [3, 1, 2, 3, 4], 1: [2, 3, 0, 1, 6]},
),
# Test case 7: List of Strings
(
["YOLO2", "YOLO3", "YOLO0", "YOLO1", "YOLO6"],
{0: ["YOLO2", "YOLO3", "YOLO0", "YOLO1", "YOLO6"]},
),
# Test case 8: List of Numbers
(
[1, 0.42, 12321, 1293, 0.00123],
{0: [1, 0.42, 12321, 1293, 0.00123]},
),
# Test case 9: Dict of Dicts
(
{
"0": {"x": 0, "y": 2},
"1": {"x": 1, "y": 3},
"2": {"x": 2, "y": 0},
"3": {"x": 3, "y": 1},
"4": {"x": 4, "y": 6},
},
{"x": [0, 1, 2, 3, 4], "y": [2, 3, 0, 1, 6]},
),
],
)
def test_to_dict(results, expected):
def test_to_dict(results, expected_dict):
"""Test helper."""
# Arrange
if results:
results = [Data(**d) for d in results]
co = OBBject(results=results)
# Act and Assert
if isinstance(expected, Exception):
with pytest.raises(expected.__class__) as exc_info:
if isinstance(expected_dict, (list, dict)):
result = co.to_dict()
assert result == expected_dict
else:
with pytest.raises(expected_dict.__class__) as exc_info:
co.to_dict()
assert str(exc_info.value) == str(expected)
else:
result = co.to_dict()
result.pop("index", None)
assert result == expected
assert str(exc_info.value) == str(expected_dict)
@patch("openbb_core.app.model.obbject.OBBject.to_dataframe")