Parameterize obbject.to_df with index and sort_by (#5877)

* Parameterize obbject.to_df with index and sort_by

* unit tests for obbject

* keep using the date as default to use on  basemode_to_df()

---------

Co-authored-by: Igor Radovanovic <74266147+IgorWounds@users.noreply.github.com>

openbb_platform/core/openbb_core/app/model/obbject.py

@@ -11,6 +11,7 @@ from typing import (
     Optional,
     Set,
     TypeVar,
+    Union,
 )
 
 import pandas as pd
@@ -24,6 +25,7 @@ from openbb_core.app.model.charts.chart import Chart
 from openbb_core.app.provider_interface import ProviderInterface
 from openbb_core.app.query import Query
 from openbb_core.app.utils import basemodel_to_df
+from openbb_core.provider.abstract.data import Data
 
 if TYPE_CHECKING:
     try:
@@ -89,11 +91,15 @@ class OBBject(Tagged, Generic[T]):
         """Return the model name with the parameters."""
         return f"OBBject[{cls.results_type_repr(params)}]"
 
-    def to_df(self) -> pd.DataFrame:
+    def to_df(
+        self, index: Optional[str] = None, sort_by: Optional[str] = None
+    ) -> pd.DataFrame:
         """Alias for `to_dataframe`."""
-        return self.to_dataframe()
+        return self.to_dataframe(index=index, sort_by=sort_by)
 
-    def to_dataframe(self) -> pd.DataFrame:
+    def to_dataframe(
+        self, index: Optional[str] = None, sort_by: Optional[str] = None
+    ) -> pd.DataFrame:
         """Convert results field to pandas dataframe.
 
         Supports converting creating pandas DataFrames from the following
@@ -109,18 +115,25 @@ class OBBject(Tagged, Generic[T]):
         - Dict[str, List]
         - Dict[str, BaseModel]
 
+        Parameters
+        ----------
+        index : Optional[str]
+            Column name to use as index.
+        sort_by : Optional[str]
+            Column name to sort by.
+
         Returns
         -------
         pd.DataFrame
             Pandas dataframe.
         """
 
-        def is_list_of_basemodel(items: List[Any]) -> bool:
+        def is_list_of_basemodel(items: Union[List[T], T]) -> bool:
             return isinstance(items, list) and all(
                 isinstance(item, BaseModel) for item in items
             )
 
-        if self.results is None or self.results == []:
+        if self.results is None or not self.results:
             raise OpenBBError("Results not found.")
 
         if isinstance(self.results, pd.DataFrame):
@@ -139,7 +152,7 @@ class OBBject(Tagged, Generic[T]):
                 for k, v in r.items():
                     # Dict[str, List[BaseModel]]
                     if is_list_of_basemodel(v):
-                        dict_of_df[k] = basemodel_to_df(v, "date")
+                        dict_of_df[k] = basemodel_to_df(v, index or "date")
                         sort_columns = False
                     # Dict[str, Any]
                     else:
@@ -149,12 +162,17 @@ class OBBject(Tagged, Generic[T]):
 
             # List[BaseModel]
             elif is_list_of_basemodel(res):
-                df = basemodel_to_df(res, "date")
+                dt: Union[List[Data], Data] = res  # type: ignore
+                df = basemodel_to_df(dt, index or "date")
                 sort_columns = False
             # List[List | str | int | float] | Dict[str, Dict | List | BaseModel]
             else:
                 try:
                     df = pd.DataFrame(res)
+                    # Set index, if any
+                    if index and index in df.columns:
+                        df.set_index(index, inplace=True)
+
                 except ValueError:
                     if isinstance(res, dict):
                         df = pd.DataFrame([res])
@@ -167,6 +185,10 @@ class OBBject(Tagged, Generic[T]):
                 df.sort_index(axis=1, inplace=True)
             df = df.dropna(axis=1, how="all")
 
+            # Sort by specified column
+            if sort_by:
+                df.sort_values(by=sort_by, inplace=True)
+
         except OpenBBError as e:
             raise e
         except ValueError as ve:
@@ -223,7 +245,8 @@ class OBBject(Tagged, Generic[T]):
             if not isinstance(self.results, dict):
                 transpose = False
             else:  # Only enter the loop if self.results is a dictionary
-                for key, value in self.results.items():
+                self.results: Dict[str, Any] = self.results  # type: ignore
+                for _, value in self.results.items():
                     if not isinstance(value, dict):
                         transpose = False
                         break
@@ -237,6 +260,7 @@ class OBBject(Tagged, Generic[T]):
     def to_chart(self, **kwargs):
         """
         Create or update the `Chart`.
+
         This function assumes that the provided data is a time series, if it's not, it will
         most likely result in an Exception.
 
@@ -287,5 +311,4 @@ class OBBject(Tagged, Generic[T]):
         OBBject[ResultsType]
             OBBject with results.
         """
-
         return cls(results=await query.execute())

openbb_platform/core/tests/app/model/test_obbject.py

@@ -16,7 +16,7 @@ def test_OBBject():
 
 def test_fields():
     """Smoke test."""
-    fields = OBBject.__fields__.keys()
+    fields = OBBject.model_fields.keys()
 
     assert "results" in fields
     assert "provider" in fields
@@ -218,6 +218,43 @@ def test_to_dataframe(results, expected_df):
         assert str(exc_info.value) == str(expected_df)
 
 
+@pytest.mark.parametrize(
+    "results, index, sort_by",
+    [
+        # Test case 1: Normal results with "date" column
+        (
+            [{"date": "2023-07-30", "value": 10}, {"date": "2023-07-31", "value": 20}],
+            "date",
+            "value",
+        ),
+        # Test case 2: List of Data
+        (
+            [
+                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),
+            ],
+            "x",
+            "y",
+        ),
+    ],
+)
+def test_to_dataframe_w_args(results, index, sort_by):
+    """Test helper."""
+    # Arrange
+    co = OBBject(results=results)
+
+    # Act and Assert
+    result = co.to_dataframe(index=index, sort_by=sort_by)
+    assert isinstance(result, pd.DataFrame)
+    assert result.index.name == index
+
+    # check if dataframe is properly sorted
+    assert result[sort_by].is_monotonic_increasing
+
+
 @pytest.mark.parametrize(
     "results, expected_dict",
     [  # Case 1: Normal results with "date" column