diff --git a/openbb_sdk/sdk/core/openbb_core/app/model/obbject.py b/openbb_sdk/sdk/core/openbb_core/app/model/obbject.py index 405fec1eec4..9bb84def10a 100644 --- a/openbb_sdk/sdk/core/openbb_core/app/model/obbject.py +++ b/openbb_sdk/sdk/core/openbb_core/app/model/obbject.py @@ -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): diff --git a/openbb_sdk/sdk/core/openbb_core/app/utils.py b/openbb_sdk/sdk/core/openbb_core/app/utils.py index 960e63b9028..5e1874fcdd3 100644 --- a/openbb_sdk/sdk/core/openbb_core/app/utils.py +++ b/openbb_sdk/sdk/core/openbb_core/app/utils.py @@ -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 diff --git a/openbb_sdk/sdk/core/tests/app/model/test_obbject.py b/openbb_sdk/sdk/core/tests/app/model/test_obbject.py index 8a0e9ee9f9d..a6b8c745b9d 100644 --- a/openbb_sdk/sdk/core/tests/app/model/test_obbject.py +++ b/openbb_sdk/sdk/core/tests/app/model/test_obbject.py @@ -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")