From e96e53d1069797779d0d13fbb638d405b03630e9 Mon Sep 17 00:00:00 2001
From: swist <me@swistofon.pl>
Date: Wed, 23 Oct 2019 17:37:56 +0100
Subject: [PATCH] Loosen up regex requirements in resolve single

OmegaConf's templating language allows very few characters.
There are some use cases where you'd want to have special characters
so that your custom resolver can handle resolution there (kind of like
embedding a templating language inside of OmegaConfs). This allows
such behaviour without seemingly breaking anything too majorly
---
 docs/source/usage.rst       | 21 +++++++++++++++++++++
 omegaconf/config.py         |  2 +-
 omegaconf/omegaconf.py      | 17 ++++++++++++++++-
 tests/test_base_config.py   | 17 +++++++++++++++++
 tests/test_interpolation.py | 25 ++++++++++++++++++++++++-
 5 files changed, 79 insertions(+), 3 deletions(-)

diff --git a/docs/source/usage.rst b/docs/source/usage.rst
index e383efbf5..d58aba5c6 100644
--- a/docs/source/usage.rst
+++ b/docs/source/usage.rst
@@ -255,6 +255,27 @@ This example creates a resolver that adds 10 the the given value.
     1000
 
 
+Custom resolvers also support variadic argument lists. Commas designate the boundaries between arguments.
+Any you can escape them using the :code:`\` character. The strings will get stripped on both ends unles you escape
+the whitespace with :code:`\`
+
+.. doctest::
+
+    >>> OmegaConf.register_resolver("concat", lambda x,y: x+y)
+    >>> c = OmegaConf.create({
+    ...     'key': '${concat:Hello\ ,World}'
+    ...     'no_escape': '${concat:Hello, World}'
+    ...     'escape_whitespace': '${concat:Hello,\ World}'
+    ... })
+    >>> c.key
+    "Hello World"
+    >>> c.no_escape
+    "HelloWorld"
+    >>> c.escape_whitespace
+    "Hello World"
+
+
+
 Merging configurations
 ----------------------
 Merging configurations enables the creation of reusable configuration files for each logical component
diff --git a/omegaconf/config.py b/omegaconf/config.py
index 11d5499ba..256cfef1b 100644
--- a/omegaconf/config.py
+++ b/omegaconf/config.py
@@ -416,7 +416,7 @@ def _resolve_value(root_node, inter_type, inter_key):
         return ret
 
     def _resolve_single(self, value):
-        match_list = list(re.finditer(r"\${(\w+:)?([\w\.%_-]+?)}", value))
+        match_list = list(re.finditer(r"\${(\w+:)?(.+?)?}", value))
         if len(match_list) == 0:
             return value
 
diff --git a/omegaconf/omegaconf.py b/omegaconf/omegaconf.py
index e00b6bef5..df66a2b25 100644
--- a/omegaconf/omegaconf.py
+++ b/omegaconf/omegaconf.py
@@ -2,6 +2,7 @@
 import copy
 import io
 import os
+import re
 import sys
 import warnings
 from contextlib import contextmanager
@@ -137,6 +138,20 @@ def merge(*others):
         target.merge_with(*others[1:])
         return target
 
+    @staticmethod
+    def _unescape_word_boundary(match):
+        if match.start() == 0 or match.end() == len(match.string):
+            return ''
+        return match.group(0)
+
+    @staticmethod
+    def _tokenize_args(string):
+        if string is None or string == '':
+            return []
+        escaped = re.split(r'(?<!\\),', string)
+        escaped = [re.sub(r'(?<!\\) ', OmegaConf._unescape_word_boundary, x) for x in escaped]
+        return [re.sub(r'(\\([ ,]))', lambda x: x.group(2), x) for x in escaped]
+
     @staticmethod
     def register_resolver(name, resolver):
         assert callable(resolver), "resolver must be callable"
@@ -145,7 +160,7 @@ def register_resolver(name, resolver):
 
         def caching(config, key):
             cache = OmegaConf.get_cache(config)[name]
-            val = cache[key] if key in cache else resolver(key)
+            val = cache[key] if key in cache else resolver(*OmegaConf._tokenize_args(key))
             cache[key] = val
             return val
 
diff --git a/tests/test_base_config.py b/tests/test_base_config.py
index d0eeb3980..713ff52df 100644
--- a/tests/test_base_config.py
+++ b/tests/test_base_config.py
@@ -245,6 +245,23 @@ def test_read_write_override(src, func, expectation):
             func(c)
 
 
+@pytest.mark.parametrize('string, tokenized', [
+    ('dog,cat', ['dog', 'cat']),
+    ('dog\,cat\ ', ['dog,cat ']),
+    ('dog,\ cat', ['dog', ' cat']),
+    ('\ ,cat', [' ', 'cat']),
+    ('dog, cat', ['dog', 'cat']),
+    ('dog, ca t', ['dog', 'ca t']),
+    ('dog, cat', ['dog', 'cat']),
+    ('whitespace\ , before comma', ['whitespace ', 'before comma']),
+    (None, []),
+    ('', []),
+    ('no , escape', ['no', 'escape']),
+])
+def test_tokenize_with_escapes(string, tokenized):
+    assert OmegaConf._tokenize_args(string) == tokenized
+
+
 @pytest.mark.parametrize('src, func, expectation', [
     ({}, lambda c: c.__setattr__('foo', 1), pytest.raises(KeyError)),
 ])
diff --git a/tests/test_interpolation.py b/tests/test_interpolation.py
index 6196298e5..18fc04aee 100644
--- a/tests/test_interpolation.py
+++ b/tests/test_interpolation.py
@@ -233,7 +233,7 @@ def test_resolver_cache_1():
 
 def test_resolver_cache_2():
     """
-    Tests that resolver cache is not shared between different OmegaConf objects 
+    Tests that resolver cache is not shared between different OmegaConf objects
     """
     try:
         OmegaConf.register_resolver("random", lambda _: random.randint(0, 10000000))
@@ -250,6 +250,29 @@ def test_resolver_cache_2():
         OmegaConf.clear_resolvers()
 
 
+def test_resolver_that_allows_a_list_of_arguments():
+    try:
+        OmegaConf.register_resolver('arg_list', lambda *args: args)
+        OmegaConf.register_resolver('zero_arg', lambda: 'zero')
+        os.environ["ENV_DEFAULT"] = "World"
+        c = OmegaConf.create(
+            dict(
+                arg_list="${arg_list:cat, dog}",
+                escape_comma="${arg_list:cat\, do g}",
+                escape_whitespace="${arg_list:cat,\ dog}",
+                zero_arg="${zero_arg:}",
+            )
+        )
+
+        assert c.arg_list == ("cat", "dog")
+        assert c.escape_comma == ("cat, do g",)
+        assert c.escape_whitespace == ("cat", " dog")
+        assert c.zero_arg == "zero"
+    finally:
+        del os.environ['ENV_DEFAULT']
+        OmegaConf.clear_resolvers
+
+
 def test_copy_cache():
     OmegaConf.register_resolver("random", lambda _: random.randint(0, 10000000))
     c1 = OmegaConf.create(dict(