import pandas as pd
from spotforecast2_safe.preprocessing.data_transform import expand_index
# DatetimeIndex
dates = pd.date_range("2023-01-01", periods=5, freq="D")
new_index = expand_index(dates, 3)
print(new_index)
assert len(new_index) == 3
assert str(new_index[0].date()) == "2023-01-06"
# RangeIndex
range_idx = pd.RangeIndex(start=0, stop=10)
new_index = expand_index(range_idx, 5)
print(new_index)
assert new_index.equals(pd.RangeIndex(start=10, stop=15, step=1))
# None index (creates new RangeIndex)
new_index = expand_index(None, 3)
print(new_index)
assert new_index.equals(pd.RangeIndex(start=0, stop=3, step=1))
# Invalid: steps not an integer raises TypeError
try:
expand_index(dates, 3.5)
except TypeError as e:
print("Error: steps must be an integer")DatetimeIndex(['2023-01-06', '2023-01-07', '2023-01-08'], dtype='datetime64[us]', freq='D')
RangeIndex(start=10, stop=15, step=1)
RangeIndex(start=0, stop=3, step=1)
Error: steps must be an integer