Update type hint and consolidate init

alteryx · Aug 7, 2023 · 76cf807 · 76cf807
1 parent c683c2f
commit 76cf807
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 13 deletions.
diff --git a/evalml/pipelines/components/estimators/regressors/varmax_regressor.py b/evalml/pipelines/components/estimators/regressors/varmax_regressor.py
@@ -115,7 +115,7 @@ def _set_forecast_horizon(self, X: pd.DataFrame):
         )
         return fh_
 
-    def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
+    def fit(self, X: pd.DataFrame, y: Optional[pd.DataFrame] = None):
         """Fits VARMAX regressor to data.
 
         Args:
@@ -171,7 +171,7 @@ def _manage_types_and_forecast(self, X: pd.DataFrame) -> tuple:
         )
         return X, fh_
 
-    def predict(self, X: pd.DataFrame, y: Optional[pd.Series] = None) -> pd.Series:
+    def predict(self, X: pd.DataFrame, y: Optional[pd.DataFrame] = None) -> pd.Series:
         """Make predictions using fitted VARMAX regressor.
 
         Args:

diff --git a/evalml/tests/component_tests/test_varmax_regressor.py b/evalml/tests/component_tests/test_varmax_regressor.py
@@ -205,18 +205,18 @@ def test_different_time_units_out_of_sample(
 def test_varmax_supports_boolean_features():
     from sktime.forecasting.varmax import VARMAX
 
-    X = pd.DataFrame({"dates": pd.date_range("2021-01-01", periods=10)})
-    X.ww.init()
-    X.ww["bool_1"] = (
-        pd.Series([True, False])
-        .sample(n=10, replace=True, random_state=0)
-        .reset_index(drop=True)
-    )
-    X.ww["bool_2"] = (
-        pd.Series([True, False])
-        .sample(n=10, replace=True, random_state=1)
-        .reset_index(drop=True)
+    X = pd.DataFrame(
+        {
+            "dates": pd.date_range("2021-01-01", periods=10),
+            "bool_1": pd.Series([True, False])
+            .sample(n=10, replace=True, random_state=0)
+            .reset_index(drop=True),
+            "bool_2": pd.Series([True, False])
+            .sample(n=10, replace=True, random_state=1)
+            .reset_index(drop=True),
+        },
     )
+    X.ww.init()
     y = pd.DataFrame({"target_1": np.random.rand(10), "target_2": np.random.rand(10)})
 
     vx = VARMAXRegressor(time_index="dates")