Enabling passing context into Model.forecast v2

CHANGELOG.md

@@ -24,7 +24,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Changed
 - 
 - 
-- 
+- Changed hierarchy of base models, enable passing context into models ([#888](https://github.com/tinkoff-ai/etna/pull/888))
 - 
 - 
 - Teach AutoARIMAModel to work with out-sample predictions ([#830](https://github.com/tinkoff-ai/etna/pull/830))

etna/models/__init__.py

@@ -1,9 +1,13 @@
 from etna import SETTINGS
 from etna.models.autoarima import AutoARIMAModel
 from etna.models.base import BaseAdapter
-from etna.models.base import BaseModel
-from etna.models.base import Model
-from etna.models.base import PerSegmentModel
+from etna.models.base import ContextIgnorantModelType
+from etna.models.base import ContextRequiredModelType
+from etna.models.base import ModelType
+from etna.models.base import NonPredictionIntervalContextIgnorantAbstractModel
+from etna.models.base import NonPredictionIntervalContextRequiredAbstractModel
+from etna.models.base import PredictionIntervalContextIgnorantAbstractModel
+from etna.models.base import PredictionIntervalContextRequiredAbstractModel
 from etna.models.catboost import CatBoostModelMultiSegment
 from etna.models.catboost import CatBoostModelPerSegment
 from etna.models.catboost import CatBoostMultiSegmentModel

etna/models/autoarima.py

@@ -5,7 +5,9 @@
 from statsmodels.tools.sm_exceptions import ValueWarning
 from statsmodels.tsa.statespace.sarimax import SARIMAXResultsWrapper
 
-from etna.models.base import PerSegmentPredictionIntervalModel
+from etna.models.base import PerSegmentModelMixin
+from etna.models.base import PredictionIntervalContextIgnorantAbstractModel
+from etna.models.base import PredictionIntervalContextIgnorantModelMixin
 from etna.models.sarimax import _SARIMAXBaseAdapter
 
 warnings.filterwarnings(
@@ -48,7 +50,9 @@ def _get_fit_results(self, endog: pd.Series, exog: pd.DataFrame) -> SARIMAXResul
         return model.arima_res_
 
 
-class AutoARIMAModel(PerSegmentPredictionIntervalModel):
+class AutoARIMAModel(
+    PerSegmentModelMixin, PredictionIntervalContextIgnorantModelMixin, PredictionIntervalContextIgnorantAbstractModel
+):
     """
     Class for holding auto arima model.
 
@@ -57,8 +61,6 @@ class AutoARIMAModel(PerSegmentPredictionIntervalModel):
     We use :py:class:`pmdarima.arima.arima.ARIMA`.
     """
 
-    context_size = 0
-
     def __init__(
         self,
         **kwargs,