Fixed comments and typings

Signed-off-by: Stefano Savare <[email protected]>

gtime/model_selection/horizon_shift.py

@@ -37,6 +37,13 @@ def horizon_shift(
     2020-01-03  3.0  4.0
     2020-01-04  4.0  NaN
     2020-01-05  NaN  NaN
+    >>> horizon_shift(X, horizon=[2])
+                y_2
+    2020-01-01  2.0
+    2020-01-02  3.0
+    2020-01-03  4.0
+    2020-01-04  NaN
+    2020-01-05  NaN
 
     """
     horizon = range(1, horizon + 1) if isinstance(horizon, (int, float)) else horizon

gtime/time_series_models/ar.py

@@ -1,3 +1,5 @@
+from typing import List, Union
+
 import numpy as np
 from sklearn.compose import make_column_selector
 from sklearn.linear_model import LinearRegression
@@ -34,7 +36,7 @@ class AR(TimeSeriesForecastingModel):
     2000-01-01 00:00:19 -0.107707  0.052031 -0.105526
     """
 
-    def __init__(self, p: int, horizon: int):
+    def __init__(self, p: int, horizon: Union[int, List[int]]):
         features = [
             tuple((f"s{i}", Shift(i), make_column_selector(dtype_include=np.number)))
             for i in range(1, p + 1)

gtime/time_series_models/base.py

@@ -1,4 +1,4 @@
-from typing import List, Tuple
+from typing import List, Tuple, Union
 
 import sklearn
 from sklearn.base import BaseEstimator, RegressorMixin
@@ -26,7 +26,7 @@ class TimeSeriesForecastingModel(BaseEstimator, RegressorMixin):
         input of class FeatureCreation, which inherits from
         ``sklearn.compose.ColumnTransformer``. It is
         used internally to instantiate FeatureCreation
-    horizon : int, required
+    horizon : Union[int, List[int], required
         how many steps to predict in the future
     model: RegressorMixin, required
         forecasting model used for predictions
@@ -57,7 +57,7 @@ class TimeSeriesForecastingModel(BaseEstimator, RegressorMixin):
     def __init__(
         self,
         features: List[Tuple],
-        horizon: int,
+        horizon: Union[int, List[int]],
         model: RegressorMixin,
         cache_features: bool = False,
     ):