Fix MultiIndex assertion

tradingstrategy-ai · May 16, 2024 · 8419da8 · 8419da8
1 parent ec10032
commit 8419da8
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/tradingstrategy/charting/candle_chart.py b/tradingstrategy/charting/candle_chart.py
@@ -50,7 +50,7 @@ def validate_ohclv_dataframe(candles: pd.DataFrame):
     if isinstance(candles.index, pd.MultiIndex):
         levels = candles.index.levels
         # (pair_id, timestamp) index - sharing index for multiple pairs in the same dataframe
-        assert levels[0].dtype == "int64" and isinstance(levels[1], pd.DatetimeIndex), f"Bad OHLCV MultiIndex - not (pair, timestamp) data: {levels[0]} and {levels[1]}"
+        assert pd.api.types.is_integer_dtype(levels[0]) and isinstance(levels[1], pd.DatetimeIndex), f"Bad OHLCV MultiIndex - not (pair, timestamp) data: {levels[0]} and {levels[1]}"
     elif not isinstance(candles.index, pd.DatetimeIndex):
         if (
             candles.index.name not in {"date", "timestamp"} and \

diff --git a/tradingstrategy/utils/df_index.py b/tradingstrategy/utils/df_index.py
@@ -22,7 +22,7 @@ def flatten_dataframe_datetime_index(df: pd.DataFrame) -> pd.DataFrame:
     assert isinstance(df.index, pd.MultiIndex), f"Got wrong index: {type(df.index)}"
 
     # (pair id, timestamp) tuples
-    assert df.index.levels[0].dtype == "int64"
+    assert pd.api.types.is_integer_dtype(df.index.levels[0])
     assert isinstance(df.index.levels[1], pd.DatetimeIndex)
 
     new_index = df.index.get_level_values(1)  # assume pair id, timestamp tuples
@@ -49,7 +49,7 @@ def get_timestamp_index(df: pd.DataFrame) -> pd.DatetimeIndex:
     assert isinstance(df.index, pd.MultiIndex), f"Got wrong index: {type(df.index)}"
 
     # (pair id, timestamp) tuples
-    assert df.index.levels[0].dtype == "int64"
+    assert pd.api.types.is_integer_dtype(df.index.levels[0])
     assert isinstance(df.index.levels[1], pd.DatetimeIndex)
 
     return df.index.get_level_values(1)