Add concat resolver for OmegaConf

CHANGELOG.md

@@ -9,7 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 - Add `ClearMLLogger` ([#568](https://github.com/etna-team/etna/pull/568))
 - **Breaking:** Add a restriction on equal feature sets across segments in the dataset ([#581](https://github.com/etna-team/etna/pull/581))
-- 
+- Add `concat` resolver for `OmegaConf` ([#604](https://github.com/etna-team/etna/pull/604))
 - 
 - 
 - 

etna/commands/__init__.py

@@ -4,9 +4,11 @@
 
 from etna.commands.backtest_command import backtest
 from etna.commands.forecast_command import forecast
+from etna.commands.resolvers import concat
 from etna.commands.resolvers import mult
 from etna.commands.resolvers import shift
 
 OmegaConf.register_new_resolver("shift", shift)
 OmegaConf.register_new_resolver("mult", mult)
+OmegaConf.register_new_resolver("concat", concat)
 OmegaConf.register_new_resolver("sum", lambda x, y: x + y)

etna/commands/resolvers.py

@@ -1,4 +1,5 @@
 from typing import List
+from typing import Union
 
 
 def shift(steps_number: int, values: List[int]) -> List[int]:
@@ -9,3 +10,8 @@
 def mult(first: int, second: int) -> int:
     """Multiply arguments."""
     return first * second
+
+
+def concat(first: Union[int, str], second: Union[int, str]) -> str:
+    """Concat arguments."""
+    return str(first) + str(second)

tests/test_commands/test_resolvers.py

@@ -0,0 +1,78 @@
+import pytest
+from hydra_slayer import get_from_params
+
+from etna.auto.pool.utils import fill_template
+
+
+def construct_pipeline(config, horizon=7):
+    config = fill_template(config, {"horizon": horizon})
+    pipeline = get_from_params(**config)
+    return pipeline
+
+
+@pytest.fixture
+def config_with_shift():
+    return {
+        "_target_": "etna.pipeline.Pipeline",
+        "horizon": "${__aux__.horizon}",
+        "model": {"_target_": "etna.models.LinearPerSegmentModel"},
+        "transforms": [
+            {"_target_": "etna.transforms.LagTransform", "in_column": "target", "lags": "${shift:${horizon},[1, 2]}"}
+        ],
+    }
+
+
+@pytest.fixture
+def config_with_mult():
+    return {
+        "_target_": "etna.pipeline.Pipeline",
+        "horizon": "${__aux__.horizon}",
+        "model": {"_target_": "etna.models.LinearPerSegmentModel"},
+        "transforms": [
+            {"_target_": "etna.transforms.LagTransform", "in_column": "target", "lags": ["${mult:${horizon},2}"]}
+        ],
+    }
+
+
+@pytest.fixture
+def config_with_concat():
+    return {
+        "_target_": "etna.pipeline.Pipeline",
+        "horizon": "${__aux__.horizon}",
+        "model": {"_target_": "etna.models.LinearPerSegmentModel"},
+        "transforms": [
+            {"_target_": "etna.transforms.LagTransform", "in_column": "${concat:tar,get}", "lags": ["${horizon}"]}
+        ],
+    }
+
+
+@pytest.fixture
+def config_with_sum():
+    return {
+        "_target_": "etna.pipeline.Pipeline",
+        "horizon": "${__aux__.horizon}",
+        "model": {"_target_": "etna.models.LinearPerSegmentModel"},
+        "transforms": [
+            {"_target_": "etna.transforms.LagTransform", "in_column": "target", "lags": ["${sum:${horizon},2}"]}
+        ],
+    }
+
+
+def test_shift(config_with_shift):
+    pipeline = construct_pipeline(config_with_shift)
+    assert pipeline.transforms[0].lags == [8, 9]
+
+
+def test_mult(config_with_mult):
+    pipeline = construct_pipeline(config_with_mult)
+    assert pipeline.transforms[0].lags == [14]
+
+
+def test_concat(config_with_concat):
+    pipeline = construct_pipeline(config_with_concat)
+    assert pipeline.transforms[0].in_column == "target"
+
+
+def test_sum(config_with_sum):
+    pipeline = construct_pipeline(config_with_sum)
+    assert pipeline.transforms[0].lags == [9]