Merge pull request #5398 from OpenBB-finance/feature/obbject-representations

Make OBBject.to_dataframe more flexible

openbb_sdk/CONTRIBUTING.md

@@ -420,7 +420,7 @@ Refer to the [`to_chart` implementation](openbb_sdk/extensions/charting/openbb_c
 In order to contribute to the SDK you need to setup your environment in order to ensure a smooth development experience.
 
 <details>
-<summary>Need help seting up Miniconda or Git?</summary>
+<summary>Need help setting up Miniconda or Git?</summary>
 
 Sometimes installing Miniconda or Git can be a bit tricky, so we've prepared a set of instructions to help you get started.
 

openbb_sdk/extensions/economy/openbb_economy/economy_router.py

@@ -1,3 +1,4 @@
+"""Economy Router."""
 from openbb_core.app.model.command_context import CommandContext
 from openbb_core.app.model.obbject import OBBject
 from openbb_core.app.provider_interface import (

openbb_sdk/providers/fred/openbb_fred/utils/fred_base.py

@@ -1,3 +1,4 @@
+"""Base class for Fred API."""
 from datetime import date
 from typing import Any, Optional
 from urllib.parse import urlencode
@@ -8,7 +9,16 @@ ROOT_URL = "https://api.stlouisfed.org/fred"
 
 
 class Fred:
+    """Base class for Fred API."""
+
     def __init__(self, api_key: Optional[str]):
+        """Initialize Fred class.
+
+        Parameters
+        ----------
+        api_key : str
+            API key for FRED
+        """
         self.api_key = api_key
 
     def __fetch_data(self, url: str, **kwargs: Any):
@@ -23,8 +33,9 @@ class Fred:
         end_date: Optional[date] = None,
         **kwargs,
     ) -> dict:
-        """
-        Get data for a Fred series id. This fetches the latest known data.
+        """Get data for a Fred series id.
+
+        This fetches the latest known data.
         Code copied from: https://github.com/mortada/fredapi/blob/master/fredapi/fred.py
 
         Parameters
@@ -57,4 +68,10 @@ class Fred:
         root = self.__fetch_data(url, **kwargs)
         if root is None:
             raise ValueError("No data exists for series id: " + series_id)
+        if (
+            "error_code" in root
+            and "error_message" in root
+            and root["error_message"] != ""
+        ):
+            raise ValueError(root["error_message"])
         return root["observations"]

openbb_sdk/sdk/core/openbb_core/app/model/obbject.py

@@ -1,7 +1,8 @@
-from typing import Dict, Generic, List, Literal, Optional, TypeVar, Union
+"""The OBBject."""
+from typing import Dict, Generic, List, Literal, Optional, TypeVar
 
 import pandas as pd
-from pydantic import Field
+from pydantic import BaseModel, Field
 from pydantic.generics import GenericModel
 
 from openbb_core.app.charting_service import ChartingService
@@ -42,44 +43,70 @@ class OBBject(GenericModel, Generic[T], Tagged):
     )
 
     def __repr__(self) -> str:
-        return f"{self.__class__.__name__}\n\n" + "\n".join(
-            f"{k}: {v}" for k, v in self.dict().items()
+        """Human readable representation of the object."""
+        return (
+            self.__class__.__name__
+            + "\n\n"
+            + "\n".join(
+                [
+                    f"{k}: {v}"[:83]
+                    + ("..." if len(f"{k}: {v}") > len(f"{k}: {v}"[:83]) else "")
+                    for k, v in self.dict().items()
+                ]
+            )
         )
 
-    def to_dataframe(
-        self, concat: bool = True
-    ) -> Union[pd.DataFrame, Dict[str, pd.DataFrame]]:
-        """Converts results field to pandas dataframe.
+    def to_dataframe(self) -> pd.DataFrame:
+        """Convert results field to pandas dataframe.
 
-        Parameters
-        ----------
-        concat : bool, optional
-            If True, it concatenates the dataframes, by default True.
+        Supports converting creating pandas DataFrames from the following
+        serializable data formats:
+
+        - List[BaseModel]
+        - List[Dict]
+        - List[List]
+        - List[str]
+        - List[int]
+        - List[float]
+        - Dict[str, Dict]
+        - Dict[str, List]
+        - Dict[str, BaseModel]
 
         Returns
         -------
-        Union[pd.DataFrame, Dict[str, pd.DataFrame]]
-            Pandas dataframe or dictionary of dataframes.
+        pd.DataFrame
+            Pandas dataframe.
         """
-        if not self.results:
+        if self.results is None or self.results == []:
             raise OpenBBError("Results not found.")
 
         try:
             res = self.results
-            if isinstance(res, list):
-                if isinstance(res[0], dict):
-                    for r in res:
-                        dict_of_df = {
-                            k: basemodel_to_df(v, "date") for k, v in r.items()
-                        }
-                        df = pd.concat(dict_of_df, axis=1) if concat else dict_of_df
 
-                else:
-                    df = basemodel_to_df(res, "date")  # type: ignore
+            df = None
+            if isinstance(res, list) and len(res) == 1 and isinstance(res[0], dict):
+                for r in res:
+                    dict_of_df = {}
+                    for k, v in r.items():
+                        if isinstance(v, list) and all(
+                            isinstance(nv, BaseModel) for nv in v
+                        ):
+                            dict_of_df[k] = basemodel_to_df(v, "date")
+                        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")
             else:
-                df = basemodel_to_df(res, "date")  # type: ignore
+                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)
             df = df.dropna(axis=1, how="all")
 
         except Exception as e:
@@ -88,7 +115,7 @@ class OBBject(GenericModel, Generic[T], Tagged):
         return df
 
     def to_dict(self) -> Dict[str, List]:
-        """Converts results field to list of values.
+        """Convert results field to list of values.
 
         Returns
         -------
@@ -103,8 +130,8 @@ class OBBject(GenericModel, Generic[T], Tagged):
         return results
 
     def to_chart(self, **kwargs):
-        """
-        Create or update the `Chart`.
+        """Create or update the `Chart`.
+
         Note that the `chart` attribute is composed by: `content`, `format` and `fig`.
 
         Parameters
@@ -127,8 +154,7 @@ class OBBject(GenericModel, Generic[T], Tagged):
         return self.chart.fig
 
     def show(self):
-        """Displays chart."""
-
+        """Display chart."""
         if not self.chart or not self.chart.fig:
             raise OpenBBError(
                 "Chart not found. "

openbb_sdk/sdk/core/tests/app/model/test_obbject.py

@@ -1,3 +1,4 @@
+"""Tests for the OBBject class."""
 from unittest.mock import MagicMock, patch
 
 import pandas as pd
@@ -7,11 +8,13 @@ from openbb_provider.abstract.data import Data
 
 
 def test_OBBject():
+    """Smoke test."""
     co = OBBject()
     assert isinstance(co, OBBject)
 
 
 def test_fields():
+    """Smoke test."""
     fields = OBBject.__fields__.keys()
 
     assert "results" in fields
@@ -21,6 +24,7 @@ def test_fields():
 
 
 def test_to_dataframe_no_results():
+    """Test helper."""
     co = OBBject()
     with pytest.raises(Exception):
         co.to_dataframe()
@@ -33,8 +37,10 @@ def test_to_dataframe_no_results():
         (
             [{"date": "2023-07-30", "value": 10}, {"date": "2023-07-31", "value": 20}],
             pd.DataFrame(
-                {"value": [10, 20]},
-                index=pd.to_datetime(["2023-07-30", "2023-07-31"]),
+                [
+                    {"date": "2023-07-30", "value": 10},
+                    {"date": "2023-07-31", "value": 20},
+                ],
             ),
         ),
         # Test case 2: Normal results without "date" column
@@ -42,16 +48,135 @@ def test_to_dataframe_no_results():
             [{"value": 10}, {"value": 20}],
             pd.DataFrame({"value": [10, 20]}, index=pd.RangeIndex(start=0, stop=2)),
         ),
-        # Test case 3: Empty 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),
+            ],
+            pd.DataFrame({"x": [0, 1, 2, 3, 4], "y": [2, 3, 0, 1, 6]}),
+        ),
+        # 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},
+            ],
+            pd.DataFrame({"x": [1, 1, 2, 3, 4], "y": [2, 3, 0, 1, 6]}),
+        ),
+        # Test case 5: List of Lists
+        (
+            [[0, 1], [1, 3], [2, 0], [3, 1], [4, 6]],
+            pd.DataFrame([[0, 1], [1, 3], [2, 0], [3, 1], [4, 6]]),
+        ),
+        # Test case 6: List of Tuples
+        (
+            [(3, 2), (1, 3), (2, 0), (3, 1), (4, 6)],
+            pd.DataFrame([(3, 2), (1, 3), (2, 0), (3, 1), (4, 6)]),
+        ),
+        # Test case 7: List of Strings
+        (
+            ["YOLO2", "YOLO3", "YOLO0", "YOLO1", "YOLO6"],
+            pd.DataFrame(["YOLO2", "YOLO3", "YOLO0", "YOLO1", "YOLO6"]),
+        ),
+        # Test case 7: List of Numbers
+        (
+            [1, 0.42, 12321298, 129387129387192837, 0.000000123],
+            pd.DataFrame([1, 0.42, 12321298, 129387129387192837, 0.000000123]),
+        ),
+        # Test case 7: 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},
+            },
+            pd.DataFrame(
+                {
+                    "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},
+                },
+            ),
+        ),
+        # Test case 8: Dict of Lists
+        (
+            {"0": [0, 2], "1": [1, 3], "2": [2, 0], "3": [3, 1], "4": [4, 6]},
+            pd.DataFrame(
+                {"0": [0, 2], "1": [1, 3], "2": [2, 0], "3": [3, 1], "4": [4, 6]}
+            ),
+        ),
+        # Test case 9: List of dict of data
+        (
+            [
+                {
+                    "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),
+                    ],
+                    "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),
+                    ],
+                }
+            ],
+            pd.concat(
+                {
+                    "df1": pd.DataFrame(
+                        {
+                            "date": [
+                                pd.to_datetime("1956-01-01"),
+                                pd.to_datetime("1956-02-01"),
+                                pd.to_datetime("1956-03-01"),
+                            ],
+                            "another_date": [
+                                "2023-09-01",
+                                "2023-09-01",
+                                "2023-09-01",
+                            ],
+                            "value": [0.0, 0.0, 0.0],
+                        }
+                    ).set_index("date"),
+                    "df2": pd.DataFrame(
+                        {
+                            "date": [
+                                pd.to_datetime("1955-03-01"),
+                                pd.to_datetime("1955-04-01"),
+                                pd.to_datetime("1955-05-01"),
+                            ],
+                            "another_date": [
+                                "2023-09-01",
+                                "2023-09-01",
+                                "2023-09-01",
+                            ],
+                            "value": [0.0, 0.0, 0.0],
+                        }
+                    ).set_index("date"),
+                },
+                axis=1,
+            ),
+        ),
+        # Test case 10: Empty results
         ([], OpenBBError("Results not found.")),
-        # Test case 4: Results as None, should raise OpenBBError
+        # Test case 11: Results as None, should raise OpenBBError
         (None, OpenBBError("Results not found.")),
     ],
 )
 def test_to_dataframe(results, expected_df):
+    """Test helper."""
     # Arrange
-    if results:
-        results = [Data(**d) for d in results]
     co = OBBject(results=results)
 
     # Act and Assert
@@ -88,6 +213,7 @@ def test_to_dataframe(results, expected_df):
     ],
 )
 def test_to_dict(results, expected):
+    """Test helper."""
     # Arrange
     if results:
         results = [Data(**d) for d in results]
@@ -112,6 +238,8 @@ def test_to_chart_with_new_chart(
     mock_charting_service,
     mock_to_dataframe,
 ):
+    """Test helper."""
+
     def get_mock_dataframe():
         data = {
             "col1": [1, 2, 3],
@@ -144,6 +272,7 @@ def test_to_chart_with_new_chart(
 
 
 def test_show_chart_exists():
+    """Test helper."""
     mock_instance = OBBject()
     # Arrange
     mock_instance.chart = MagicMock(spec=Chart)
@@ -158,6 +287,7 @@ def test_show_chart_exists():
 
 
 def test_show_chart_no_chart():
+    """Test helper."""
     mock_instance = OBBject()
 
     # Act and Assert
@@ -166,6 +296,7 @@ def test_show_chart_no_chart():
 
 
 def test_show_chart_no_fig():
+    """Test helper."""
     mock_instance = OBBject()
     # Arrange
     mock_instance.chart = Chart()

pyproject.toml

@@ -188,10 +188,14 @@ requires = ["setuptools<65.5.0", "poetry-core>=1.0.0"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.pydocstyle]
+inherit = true
 convention = "numpy"
 match = '((?!test_).)*\.py'
 add-ignore = "D401"
 
+[tool.flake8]
+max-line-length = 122
+ignore = "S101"
 
 [tool.ruff]
 line-length = 122