Repository: kedro-org/kedro

Issue:
Allow registering of custom resolvers to OmegaConfigLoader through `CONFIG_LOADER_ARGS`
## Description
Allowing users to register custom resolvers will make it easy for users to do more advanced interpolation in their configs.
OmegaConf documentation about custom resolvers - https://omegaconf.readthedocs.io/en/2.3_branch/custom_resolvers.html

## Context
Follow up task for #2407

## Possible Implementation
Users should be able to add custom resolvers in their `settings.py` files -
```python
from kedro.config import OmegaConfigLoader


CONFIG_LOADER_CLASS = OmegaConfigLoader

CONFIG_LOADER_ARGS = {
    "custom_resolvers": {
        "add": lambda *x: sum(x),
    }
}
```
The config files then could look like this -
`parameters.yml`
```yaml
model_options:
  test_size: ${add:1,2,3}
  random_state: 3
```

* The implementation of this on the `OmegaConfigLoader` side should be a very thin wrapper for `OmegaConf.register_new_resolver()` so that any future changes to `OmegaConf` don't affect anything on Kedro side.
* Users can overwrite built-in resolvers using `replace=True`



Discussion:
ankatiyar: I've added a note to #2624 to document the fact that `replace` flag is set to True
ankatiyar: @merelcht, I've turned off "replace=True" after discussion with @noklam. The only use case for that would be overwriting the built-in resolvers (oc.create, oc.decode, oc.deprecate, oc.select, oc.dict) which probably not a lot of people would want to do. To switch on `oc.env` does not require it being replaced since we clear it entirely at initialisation of `OmegaConfigLoader`. `OmegaConf` also does not allow replacing resolvers by default. Either way, it doesn't make too much of a difference to the implementation. What do you think?

Also updated the test to check if registering `oc.env` works properly.
ankatiyar: For advanced users that want to overwrite the built-in resolvers, they can always do so directly in `settings.py` -
```python
from omegaconf import OmegaConf
OmegaConf.register_new_resolver(name, resolver, replace=True)
```

noklam: > > @merelcht, I've turned off "replace=True" after discussion with @noklam. The only use case for that would be overwriting the built-in resolvers (oc.create, oc.decode, oc.deprecate, oc.select, oc.dict) which probably not a lot of people would want to do. To switch on `oc.env` does not require it being replaced since we clear it entirely at initialisation of `OmegaConfigLoader`. `OmegaConf` also does not allow replacing resolvers by default. Either way, it doesn't make too much of a difference to the implementation. What do you think?
> 
> I think that's a fair point. If it turns out that users want to be able to replace resolvers we could potentially add another argument to `CONFIG_LOADER_ARGS` where they can overwrite it, e.g. `replace_resolvers`. Thanks for discussing this in more detail. I'm happy with this approach 🙂 
>  
> > Also updated the test to check if registering `oc.env` works properly.
> 
> 👍 
> 
> 

@merelcht We also discussed this point. Adding a new replace argument isn't hard, but it  clutters the already long argument list.

Two reasons why I think we shouldn't do it now.
1. If we want to add replace argument, arguably we may need to add a few more since omegaconf allows few extra options. If so should we provide a omegaconf_kwargs instead of just one argument? It's not clear for now.
2. We can always add a new one, and it wouldn't be a breaking change. However if we decide to add replace as an arguments, but later we want to change it to kwargs, it will be a breaking change and we cannot do it easily.

So for now I think we have a good workaround for advance use case which won't be blocking. If there are more evidence coming in later we can then discuss more about 1. and have a hollistic view of what other arguments we may need to support. 
noklam: I also open a new issue to fix the tests and signature in one go. Will refine the description later. 

Patch:
diff --git a/RELEASE.md b/RELEASE.md
index 2fcf553b..ae0714b8 100644
--- a/RELEASE.md
+++ b/RELEASE.md
@@ -11,6 +11,7 @@
 # Upcoming Release 0.18.13
 
 ## Major features and improvements
+* Allowed registering of custom resolvers to `OmegaConfigLoader` through `CONFIG_LOADER_ARGS`.
 
 ## Bug fixes and other changes
 
diff --git a/kedro/config/omegaconf_config.py b/kedro/config/omegaconf_config.py
index d7d9bd24..4d2ace59 100644
--- a/kedro/config/omegaconf_config.py
+++ b/kedro/config/omegaconf_config.py
@@ -7,7 +7,7 @@ import io
 import logging
 import mimetypes
 from pathlib import Path
-from typing import Any, Iterable
+from typing import Any, Callable, Iterable
 
 import fsspec
 from omegaconf import OmegaConf
@@ -82,6 +82,7 @@ class OmegaConfigLoader(AbstractConfigLoader):
         config_patterns: dict[str, list[str]] = None,
         base_env: str = "base",
         default_run_env: str = "local",
+        custom_resolvers: dict[str, Callable] = None,
     ):
         """Instantiates a ``OmegaConfigLoader``.
 
@@ -97,6 +98,8 @@ class OmegaConfigLoader(AbstractConfigLoader):
                 the configuration paths.
             default_run_env: Name of the default run environment. Defaults to `"local"`.
                 Can be overridden by supplying the `env` argument.
+            custom_resolvers: A dictionary of custom resolvers to be registered. For more information,
+             see here: https://omegaconf.readthedocs.io/en/2.3_branch/custom_resolvers.html#custom-resolvers
         """
         self.base_env = base_env
         self.default_run_env = default_run_env
@@ -111,6 +114,9 @@ class OmegaConfigLoader(AbstractConfigLoader):
 
         # Deactivate oc.env built-in resolver for OmegaConf
         OmegaConf.clear_resolver("oc.env")
+        # Register user provided custom resolvers
+        if custom_resolvers:
+            self._register_new_resolvers(custom_resolvers)
 
         file_mimetype, _ = mimetypes.guess_type(conf_source)
         if file_mimetype == "application/x-tar":
@@ -302,6 +308,15 @@ class OmegaConfigLoader(AbstractConfigLoader):
             ".json",
         ]
 
+    @staticmethod
+    def _register_new_resolvers(resolvers: dict[str, Callable]):
+        """Register custom resolvers"""
+        for name, resolver in resolvers.items():
+            if not OmegaConf.has_resolver(name):
+                msg = f"Registering new custom resolver: {name}"
+                _config_logger.debug(msg)
+                OmegaConf.register_new_resolver(name=name, resolver=resolver)
+
     @staticmethod
     def _check_duplicates(seen_files_to_keys: dict[Path, set[Any]]):
         duplicates = []


Test patch:
diff --git a/tests/config/test_omegaconf_config.py b/tests/config/test_omegaconf_config.py
index dd492920..af57b522 100644
--- a/tests/config/test_omegaconf_config.py
+++ b/tests/config/test_omegaconf_config.py
@@ -649,3 +649,25 @@ class TestOmegaConfigLoader:
         conf = OmegaConfigLoader(str(tmp_path))
         conf.default_run_env = ""
         assert conf["catalog"]["companies"]["type"] == "pandas.CSVDataSet"
+
+    def test_custom_resolvers(self, tmp_path):
+        base_params = tmp_path / _BASE_ENV / "parameters.yml"
+        param_config = {
+            "model_options": {
+                "param1": "${add: 3, 4}",
+                "param2": "${plus_2: 1}",
+                "param3": "${oc.env: VAR}",
+            }
+        }
+        _write_yaml(base_params, param_config)
+        custom_resolvers = {
+            "add": lambda *x: sum(x),
+            "plus_2": lambda x: x + 2,
+            "oc.env": oc.env,
+        }
+        os.environ["VAR"] = "my_env_variable"
+        conf = OmegaConfigLoader(tmp_path, custom_resolvers=custom_resolvers)
+        conf.default_run_env = ""
+        assert conf["parameters"]["model_options"]["param1"] == 7
+        assert conf["parameters"]["model_options"]["param2"] == 3
+        assert conf["parameters"]["model_options"]["param3"] == "my_env_variable"