From 3c053c4fa1476bbeecb11842b70fcd9821ba7cd7 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       | 20 ++++++++++++++++++++
 omegaconf/config.py         |  4 +++-
 omegaconf/omegaconf.py      | 16 +++++++++++++++-
 tests/test_base_config.py   | 17 +++++++++++++++++
 tests/test_interpolation.py | 17 ++++++++++++++++-
 5 files changed, 71 insertions(+), 3 deletions(-)

diff --git a/docs/source/usage.rst b/docs/source/usage.rst
index e383efbf5..6a5270214 100644
--- a/docs/source/usage.rst
+++ b/docs/source/usage.rst
@@ -255,6 +255,26 @@ This example creates a resolver that adds 10 the the given value.
     1000
 
 
+Custom resolvers support variadic argument lists in the form of a comma separated list of zero or more values (coming in OmegaConf 1.3.1).
+Whitespaces are stripped from both ends of each value ("foo,bar" is the same as "foo, bar ").
+You can use literal commas and spaces anywhere by escaping (:code:`\,` and :code:`\ `).
+.. doctest::
+
+    >>> OmegaConf.register_resolver("concat", lambda x,y: x+y)
+    >>> c = OmegaConf.create({
+    ...     'key1': '${concat:Hello,World}',
+    ...     'key_trimmed': '${concat:Hello , World}',
+    ...     'escape_whitespace': '${concat:Hello,\ World}',
+    ... })
+    >>> c.key1
+    'HelloWorld'
+    >>> c.key_trimmed
+    '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..0a8995d40 100644
--- a/omegaconf/config.py
+++ b/omegaconf/config.py
@@ -416,7 +416,9 @@ 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))
+        key_prefix = r"\${(\w+:)?"
+        legal_characters = r"([\w\.%_ \\,-]*?)}"
+        match_list = list(re.finditer(key_prefix+legal_characters, value))
         if len(match_list) == 0:
             return value
 
diff --git a/omegaconf/omegaconf.py b/omegaconf/omegaconf.py
index e00b6bef5..03e908b70 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,19 @@ def merge(*others):
         target.merge_with(*others[1:])
         return target
 
+
+    @staticmethod
+    def _tokenize_args(string):
+        if string is None or string == '':
+            return []
+        def _unescape_word_boundary(match):
+            if match.start() == 0 or match.end() == len(match.string):
+                return ''
+            return match.group(0)
+        escaped = re.split(r'(?<!\\),', string)
+        escaped = [re.sub(r'(?<!\\) ', _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 +159,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..0c358c53e 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,21 @@ def test_resolver_cache_2():
         OmegaConf.clear_resolvers()
 
 
+@pytest.mark.parametrize('resolver,name,key,result', [
+    (lambda *args: args, 'arg_list', "${my_resolver:cat, dog}",("cat", "dog")),
+    (lambda *args: args, 'escape_comma',"${my_resolver:cat\, do g}",("cat, do g",)),
+    (lambda *args: args, 'escape_whitespace',"${my_resolver:cat\, do g}",("cat, do g",)),
+    (lambda: 'zero', 'zero_arg',"${my_resolver:}",("zero")),
+])
+def test_resolver_that_allows_a_list_of_arguments(resolver, name, key, result):
+    try:
+        OmegaConf.register_resolver('my_resolver', resolver)
+        c = OmegaConf.create({name: key})
+        assert c[name] == result
+    finally:
+        OmegaConf.clear_resolvers()
+
+
 def test_copy_cache():
     OmegaConf.register_resolver("random", lambda _: random.randint(0, 10000000))
     c1 = OmegaConf.create(dict(