Allow more types of dictionary keys in overrides grammar

hydra/_internal/grammar/utils.py

@@ -12,7 +12,7 @@
 _ESC_REGEX = re.compile(f"[{re.escape(_ESC)}]+")
 
 
-def escape(s: str) -> str:
+def escape_special_characters(s: str) -> str:
     """Escape special characters in `s`"""
     matches = _ESC_REGEX.findall(s)
     if not matches:

hydra/core/override_parser/types.py

@@ -10,7 +10,7 @@
 from omegaconf import OmegaConf
 from omegaconf._utils import is_structured_config
 
-from hydra._internal.grammar.utils import escape
+from hydra._internal.grammar.utils import escape_special_characters
 from hydra.core.config_loader import ConfigLoader
 from hydra.core.object_type import ObjectType
 from hydra.errors import HydraException
@@ -143,9 +143,9 @@ def __eq__(self, other: Any) -> Any:
             return NotImplemented
 
 
-# Ideally we would use List[ElementType] and Dict[ElementType, ElementType] but Python
-# does not seem to support recursive type definitions.
-ElementType = Union[str, int, float, bool, List[Any], Dict[Any, Any]]
+# Ideally we would use List[ElementType] and Dict[str, ElementType] but Python does not seem
+# to support recursive type definitions.
+ElementType = Union[str, int, float, bool, List[Any], Dict[str, Any]]
 ParsedElementType = Optional[Union[ElementType, QuotedString]]
 TransformerType = Callable[[ParsedElementType], Any]
 
@@ -259,8 +259,14 @@ def _convert_value(value: ParsedElementType) -> Optional[ElementType]:
         if isinstance(value, list):
             return [Override._convert_value(x) for x in value]
         elif isinstance(value, dict):
+
+            # Currently only strings are allowed as dictionary keys.
+            def check_str(k: Any) -> str:
+                assert isinstance(k, str)
+                return k
+
             return {
-                Override._convert_value(k): Override._convert_value(v)
+                check_str(Override._convert_value(k)): Override._convert_value(v)
                 for k, v in value.items()
             }
         elif isinstance(value, QuotedString):
@@ -424,7 +430,7 @@ def _get_value_element_as_str(
             )
             return "{" + s + "}"
         elif isinstance(value, str):
-            return escape(value)  # ensure special characters are escaped
+            return escape_special_characters(value)
         elif isinstance(value, (int, bool, float)):
             return str(value)
         elif is_structured_config(value):

tests/test_config_loader.py

@@ -593,38 +593,6 @@ def test_sweep_config_cache(
         monkeypatch.setenv("HOME", "/another/home/dir/")
         assert sweep_cfg.home == os.getenv("HOME")
 
-    @pytest.mark.parametrize(  # type: ignore
-        "key,expected",
-        [
-            pytest.param("id123", "id123", id="id"),
-            pytest.param("123id", "123id", id="int_plus_id"),
-            pytest.param("'quoted_single'", "quoted_single", id="quoted_single"),
-            pytest.param('"quoted_double"', "quoted_double", id="quoted_double"),
-            pytest.param("'quoted_$(){}[]'", "quoted_$(){}[]", id="quoted_misc_chars"),
-            pytest.param("a/-\\+.$%*@", "a/-\\+.$%*@", id="unquoted_misc_chars"),
-            pytest.param("white space", "white space", id="whitespace"),
-            pytest.param(
-                "\\\\\\(\\)\\[\\]\\{\\}\\:\\=\\ \\\t\\,",
-                "\\()[]{}:= \t,",
-                id="unquoted_esc",
-            ),
-        ],
-    )
-    def test_dict_key_formats(
-        self, hydra_restore_singletons: Any, path: str, key: str, expected: str
-    ) -> None:
-        """Test that we can assign dictionaries with keys that are not just IDs"""
-        config_loader = ConfigLoaderImpl(
-            config_search_path=create_config_search_path(path)
-        )
-        cfg = config_loader.load_configuration(
-            config_name="config.yaml",
-            overrides=[f"+dict={{{key}: 123}}"],
-            run_mode=RunMode.RUN,
-        )
-        assert "dict" in cfg
-        assert cfg.dict == {expected: 123}
-
 
 @pytest.mark.parametrize(  # type:ignore
     "config_file, overrides",

tests/test_internal_grammar.py

@@ -0,0 +1,27 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+from pytest import mark, param
+
+from hydra._internal.grammar.utils import escape_special_characters
+
+
+@mark.parametrize(  # type: ignore
+    ("s", "expected"),
+    [
+        param("abc", "abc", id="no_esc"),
+        param("\\", "\\\\", id="esc_backslash"),
+        param("\\\\\\", "\\\\\\\\\\\\", id="esc_backslash_x3"),
+        param("()", "\\(\\)", id="esc_parentheses"),
+        param("[]", "\\[\\]", id="esc_brackets"),
+        param("{}", "\\{\\}", id="esc_braces"),
+        param(":=,", "\\:\\=\\,", id="esc_symbols"),
+        param("  \t", "\\ \\ \\\t", id="esc_ws"),
+        param(
+            "ab\\(cd{ef}[gh]): ij,kl\t",
+            "ab\\\\\\(cd\\{ef\\}\\[gh\\]\\)\\:\\ ij\\,kl\\\t",
+            id="esc_mixed",
+        ),
+    ],
+)
+def test_escape_special_characters(s: str, expected: str) -> None:
+    escaped = escape_special_characters(s)
+    assert escaped == expected

tests/test_overrides_parser.py

@@ -1030,6 +1030,16 @@ def test_override_get_value_element_method(
         pytest.param("key={a:10,b:20}", {"a": 10, "b": 20}, id="dict"),
         pytest.param("key={a:10,b:20}", {"a": 10, "b": 20}, id="dict"),
         pytest.param("key={a:10,b:[1,2,3]}", {"a": 10, "b": [1, 2, 3]}, id="dict"),
+        pytest.param("key={123id: 0}", {"123id": 0}, id="dict_key_int_plus_id"),
+        pytest.param("key={' abc ': 0}", {" abc ": 0}, id="dict_key_quoted_single"),
+        pytest.param('key={" abc ": 0}', {" abc ": 0}, id="dict_key_quoted_double"),
+        pytest.param("key={a/-\\+.$%*@: 0}", {"a/-\\+.$%*@": 0}, id="dict_key_noquote"),
+        pytest.param("key={w s: 0}", {"w s": 0}, id="dict_key_ws"),
+        pytest.param(
+            "key={\\\\\\(\\)\\[\\]\\{\\}\\:\\=\\ \\\t\\,: 0}",
+            {"\\()[]{}:= \t,": 0},
+            id="dict_key_esc",
+        ),
     ],
 )
 def test_override_value_method(override: str, expected: str) -> None: