diff --git a/openbb_platform/core/openbb_core/app/utils/__init__.py b/openbb_platform/core/openbb_core/app/utils/__init__.py index 62486a8b8c2..a953511be3a 100644 --- a/openbb_platform/core/openbb_core/app/utils/__init__.py +++ b/openbb_platform/core/openbb_core/app/utils/__init__.py @@ -17,6 +17,13 @@ if TYPE_CHECKING: from numpy import ndarray from pandas import DataFrame, Series +_MULTIINDEX_META = ("is_multiindex", "multiindex_names", "multiindex_columns") + + +def _multiindex_columns(names: list) -> list: + """Get the column names that reset_index gives the levels of a MultiIndex.""" + return [n if n is not None else f"level_{i}" for i, n in enumerate(names)] + def basemodel_to_df( data: list[Data] | Data, @@ -25,6 +32,7 @@ def basemodel_to_df( """Convert list of BaseModel to a Pandas DataFrame.""" pd = require_optional("pandas") DataFrame, to_datetime = pd.DataFrame, pd.to_datetime # type: ignore[union-attr] + is_numeric_dtype = pd.api.types.is_numeric_dtype # type: ignore[union-attr] if isinstance(data, list): df = DataFrame( @@ -40,12 +48,20 @@ def basemodel_to_df( ) if "is_multiindex" in df.columns: - col_names = ast.literal_eval(df.multiindex_names.unique()[0]) - df = df.set_index(col_names) - df = df.drop(["is_multiindex", "multiindex_names"], axis=1) + names = ast.literal_eval(df.multiindex_names.unique()[0]) + # The level columns follow the reset_index naming unless df_to_basemodel + # had to rename them to avoid a clash, in which case it stores them. + if "multiindex_columns" in df.columns: + columns = ast.literal_eval(df.multiindex_columns.unique()[0]) + else: + columns = _multiindex_columns(names) + df = df.set_index(columns).rename_axis(names) + df = df.drop(columns=[c for c in _MULTIINDEX_META if c in df.columns]) # If the date column contains dates only, convert them to a date to avoid encoding time data. - if "date" in df.columns: + # Numeric values would be read as epoch offsets, so only text and datetimes are parsed. + parsed = None + if "date" in df.columns and not is_numeric_dtype(df["date"]): try: # Parse the column in a single call instead of once per row. parsed = to_datetime(df["date"]) @@ -53,8 +69,13 @@ def basemodel_to_df( # A column that mixes UTC offsets - a series crossing a daylight # savings boundary, for instance - has no single datetime64 # representation, so fall back to parsing each value on its own. - parsed = df["date"].apply(to_datetime) + try: + parsed = df["date"].apply(to_datetime) + except (TypeError, ValueError): + # Not a date column after all, so leave it as it is. + parsed = None + if parsed is not None: df["date"] = parsed if parsed.dtype == object: @@ -79,6 +100,7 @@ def df_to_basemodel( """Convert from a Pandas DataFrame to list of BaseModel.""" pd = require_optional("pandas") MultiIndex, Series, to_datetime = pd.MultiIndex, pd.Series, pd.to_datetime # type: ignore[union-attr] + is_numeric_dtype = pd.api.types.is_numeric_dtype # type: ignore[union-attr] is_multiindex = isinstance(df.index, MultiIndex) @@ -89,15 +111,36 @@ def df_to_basemodel( # Check if df has multiindex. If so, add the index names to the df and a boolean column if isinstance(df.index, MultiIndex): + names = list(df.index.names) + default_columns = _multiindex_columns(names) + # A level whose column name is already taken, by a column or another level, + # gets a free name instead so that reset_index does not fail. + taken = set(df.columns) | set(_MULTIINDEX_META) + columns = [] + for default in default_columns: + column = default + while column in taken: + column = f"_{column}" + taken.add(column) + columns.append(column) + df = df.rename_axis(columns).reset_index() df["is_multiindex"] = True - df["multiindex_names"] = str(df.index.names) - df = df.reset_index() + df["multiindex_names"] = str(names) + if columns != default_columns: + df["multiindex_columns"] = str(columns) # Converting to JSON will add T00:00:00.000 to all dates with no time element unless we format it as a string first. - if "date" in df.columns: - df["date"] = df["date"].apply(to_datetime) - if all(t.time() == time(0, 0) for t in df["date"]): - df["date"] = df["date"].apply(lambda x: x.date().strftime("%Y-%m-%d")) + # Numeric values would be read as epoch offsets, so only text and datetimes are parsed. + if "date" in df.columns and not is_numeric_dtype(df["date"]): + try: + dates = df["date"].apply(to_datetime) + except (TypeError, ValueError): + # Not a date column after all, so leave it as it is. + dates = None + if dates is not None: + df["date"] = dates + if all(t.time() == time(0, 0) for t in df["date"]): + df["date"] = df["date"].apply(lambda x: x.date().strftime("%Y-%m-%d")) return [ Data(**d) for d in json.loads(df.to_json(orient="records", date_format="iso")) diff --git a/openbb_platform/core/tests/app/test_utils.py b/openbb_platform/core/tests/app/test_utils.py index 9607e7c75b2..d00690eebd6 100644 --- a/openbb_platform/core/tests/app/test_utils.py +++ b/openbb_platform/core/tests/app/test_utils.py @@ -75,6 +75,59 @@ def test_basemodel_to_multiindex_df(): assert isinstance(df.index, pd.MultiIndex) +@pytest.mark.parametrize( + "names", + [["x", "y"], [None, None], ["x", None], [None, "y"]], +) +def test_multiindex_round_trip(names): + """Test that a MultiIndex survives a df -> basemodel -> df round trip.""" + original = df_multiindex.copy() + original.index.names = names + + result = basemodel_to_df(df_to_basemodel(original)) + + pd.testing.assert_frame_equal(result, original) + + +@pytest.mark.parametrize( + "names, columns", + [ + (["z", "y"], ["z"]), + ([None, "y"], ["level_0"]), + (["x", "x"], ["z"]), + (["is_multiindex", "y"], ["z"]), + ], +) +def test_multiindex_round_trip_name_clash(names, columns): + """Test that index levels named after a column or each other survive a round trip.""" + original = df_multiindex.copy() + original.index.names = names + original.columns = columns + + result = basemodel_to_df(df_to_basemodel(original)) + + pd.testing.assert_frame_equal(result, original) + + +@pytest.mark.parametrize("values", [list("abcdefghij"), list(range(10))]) +def test_multiindex_round_trip_non_date_date_level(values): + """Test that an index level named date that holds non-date values survives a round trip.""" + original = df.assign(date=values).set_index(["date", "y"]) + + result = basemodel_to_df(df_to_basemodel(original)) + + pd.testing.assert_frame_equal(result, original) + + +def test_df_to_basemodel_does_not_modify_multiindex_input(): + """Test that converting a MultiIndex DataFrame leaves the input untouched.""" + original = df.set_index(["x", "y"]) + + df_to_basemodel(original) + + pd.testing.assert_frame_equal(original, df.set_index(["x", "y"])) + + def test_get_target_column(): """Test the get_target_column helper.""" target = get_target_column(df, "x")