Dict int key type

docs/source/structured_config.rst

@@ -257,8 +257,9 @@ OmegaConf verifies at runtime that your Lists contains only values of the correc
 
 Dictionaries
 ^^^^^^^^^^^^
-Dictionaries are supported as well. Keys must be strings or enums, and values can be any of any type supported by OmegaConf
-(Any, int, float, bool, str and Enums as well as arbitrary Structured configs)
+Dictionaries are supported as well. Keys must be strings, ints or enums, and values can
+be any of any type supported by OmegaConf (Any, int, float, bool, str and Enums as well
+as arbitrary Structured configs)
 
 Misc
 ----

docs/source/usage.rst

@@ -41,26 +41,28 @@ From a dictionary
 
 .. doctest::
 
-    >>> conf = OmegaConf.create({"k" : "v", "list" : [1, {"a": "1", "b": "2"}]})
+    >>> conf = OmegaConf.create({"k" : "v", "list" : [1, {"a": "1", "b": "2", 3: "c"}]})
     >>> print(OmegaConf.to_yaml(conf))
     k: v
     list:
     - 1
     - a: '1'
       b: '2'
+      3: c
     <BLANKLINE>
 
 From a list
 ^^^^^^^^^^^
 
 .. doctest::
 
-    >>> conf = OmegaConf.create([1, {"a":10, "b": {"a":10}}])
+    >>> conf = OmegaConf.create([1, {"a":10, "b": {"a":10, 123: 456}}])
     >>> print(OmegaConf.to_yaml(conf))
     - 1
     - a: 10
       b:
         a: 10
+        123: 456
     <BLANKLINE>
 
 Tuples are supported as an valid option too.
@@ -95,6 +97,7 @@ From a YAML string
     ... list:
     ... - item1
     ... - item2
+    ... 123: 456
     ... """
     >>> conf = OmegaConf.create(s)
     >>> print(OmegaConf.to_yaml(conf))
@@ -103,6 +106,7 @@ From a YAML string
     list:
     - item1
     - item2
+    123: 456
     <BLANKLINE>
 
 From a dot-list
@@ -264,7 +268,7 @@ Save/Load YAML file
 
 .. doctest:: loaded
 
-    >>> conf = OmegaConf.create({"foo": 10, "bar": 20})
+    >>> conf = OmegaConf.create({"foo": 10, "bar": 20, 123: 456})
     >>> with tempfile.NamedTemporaryFile() as fp:
     ...     OmegaConf.save(config=conf, f=fp.name)
     ...     loaded = OmegaConf.load(fp.name)
@@ -279,7 +283,7 @@ Note that the saved file may be incompatible across different major versions of
 
 .. doctest:: loaded
 
-    >>> conf = OmegaConf.create({"foo": 10, "bar": 20})
+    >>> conf = OmegaConf.create({"foo": 10, "bar": 20, 123: 456})
     >>> with tempfile.TemporaryFile() as fp:
     ...     pickle.dump(conf, fp)
     ...     fp.flush()

omegaconf/__init__.py

@@ -1,4 +1,5 @@
 from .base import Container, Node
+from .basecontainer import DictKeyType
 from .dictconfig import DictConfig
 from .errors import (
     KeyValidationError,
@@ -39,6 +40,7 @@
     "Container",
     "ListConfig",
     "DictConfig",
+    "DictKeyType",
     "OmegaConf",
     "Resolver",
     "flag_override",

omegaconf/_utils.py

@@ -491,7 +491,9 @@ def valid_value_annotation_type(type_: Any) -> bool:
 
 
 def _valid_dict_key_annotation_type(type_: Any) -> bool:
-    return type_ is None or type_ is Any or issubclass(type_, (str, Enum))
+    from omegaconf import DictKeyType
+
+    return type_ is None or type_ is Any or issubclass(type_, DictKeyType.__args__)  # type: ignore
 
 
 def is_primitive_type(type_: Any) -> bool:

omegaconf/base.py

@@ -250,7 +250,7 @@ def __setitem__(self, key: Any, value: Any) -> None:
         ...
 
     @abstractmethod
-    def __iter__(self) -> Iterator[str]:
+    def __iter__(self) -> Iterator[Any]:
         ...
 
     @abstractmethod

omegaconf/basecontainer.py

@@ -29,6 +29,8 @@
 
 DEFAULT_VALUE_MARKER: Any = str("__DEFAULT_VALUE_MARKER__")
 
+DictKeyType = Union[str, int, Enum]
+
 
 class BaseContainer(Container, ABC):
     # static
@@ -187,7 +189,7 @@ def _to_content(
         resolve: bool,
         enum_to_str: bool = False,
         exclude_structured_configs: bool = False,
-    ) -> Union[None, Any, str, Dict[str, Any], List[Any]]:
+    ) -> Union[None, Any, str, Dict[DictKeyType, Any], List[Any]]:
         from .dictconfig import DictConfig
         from .listconfig import ListConfig
 
@@ -528,7 +530,10 @@ def assign(value_key: Any, value_to_assign: Any) -> None:
 
     @staticmethod
     def _item_eq(
-        c1: Container, k1: Union[str, int], c2: Container, k2: Union[str, int]
+        c1: Container,
+        k1: Union[DictKeyType, int],
+        c2: Container,
+        k2: Union[DictKeyType, int],
     ) -> bool:
         v1 = c1._get_node(k1)
         v2 = c2._get_node(k2)

omegaconf/dictconfig.py

@@ -9,6 +9,7 @@
     List,
     MutableMapping,
     Optional,
+    Sequence,
     Tuple,
     Type,
     Union,
@@ -32,7 +33,7 @@
     valid_value_annotation_type,
 )
 from .base import Container, ContainerMetadata, Node
-from .basecontainer import DEFAULT_VALUE_MARKER, BaseContainer
+from .basecontainer import DEFAULT_VALUE_MARKER, BaseContainer, DictKeyType
 from .errors import (
     ConfigAttributeError,
     ConfigKeyError,
@@ -47,13 +48,13 @@
 from .nodes import EnumNode, ValueNode
 
 
-class DictConfig(BaseContainer, MutableMapping[str, Any]):
+class DictConfig(BaseContainer, MutableMapping[Any, Any]):
 
     _metadata: ContainerMetadata
 
     def __init__(
         self,
-        content: Union[Dict[str, Any], Any],
+        content: Union[Dict[DictKeyType, Any], Any],
         key: Any = None,
         parent: Optional[Container] = None,
         ref_type: Union[Any, Type[Any]] = Any,
@@ -245,14 +246,12 @@ def _raise_invalid_value(
         )
         raise ValidationError(msg)
 
-    def _validate_and_normalize_key(self, key: Any) -> Union[str, Enum]:
+    def _validate_and_normalize_key(self, key: Any) -> DictKeyType:
         return self._s_validate_and_normalize_key(self._metadata.key_type, key)
 
-    def _s_validate_and_normalize_key(
-        self, key_type: Any, key: Any
-    ) -> Union[str, Enum]:
+    def _s_validate_and_normalize_key(self, key_type: Any, key: Any) -> DictKeyType:
         if key_type is Any:
-            for t in (str, Enum):
+            for t in DictKeyType.__args__:  # type: ignore
                 try:
                     return self._s_validate_and_normalize_key(key_type=t, key=key)
                 except KeyValidationError:
@@ -264,6 +263,13 @@ def _s_validate_and_normalize_key(
                     f"Key $KEY ($KEY_TYPE) is incompatible with ({key_type.__name__})"
                 )
 
+            return key
+        elif key_type == int:
+            if not isinstance(key, int):
+                raise KeyValidationError(
+                    f"Key $KEY ($KEY_TYPE) is incompatible with ({key_type.__name__})"
+                )
+
             return key
         elif issubclass(key_type, Enum):
             try:
@@ -278,7 +284,7 @@ def _s_validate_and_normalize_key(
         else:
             assert False, f"Unsupported key type {key_type}"
 
-    def __setitem__(self, key: Union[str, Enum], value: Any) -> None:
+    def __setitem__(self, key: DictKeyType, value: Any) -> None:
         try:
             self.__set_impl(key=key, value=value)
         except AttributeError as e:
@@ -288,7 +294,7 @@ def __setitem__(self, key: Union[str, Enum], value: Any) -> None:
         except Exception as e:
             self._format_and_raise(key=key, value=value, cause=e)
 
-    def __set_impl(self, key: Union[str, Enum], value: Any) -> None:
+    def __set_impl(self, key: DictKeyType, value: Any) -> None:
         key = self._validate_and_normalize_key(key)
         self._set_item_impl(key, value)
 
@@ -331,7 +337,7 @@ def __getattr__(self, key: str) -> Any:
         except Exception as e:
             self._format_and_raise(key=key, value=None, cause=e)
 
-    def __getitem__(self, key: Union[str, Enum]) -> Any:
+    def __getitem__(self, key: DictKeyType) -> Any:
         """
         Allow map style access
         :param key:
@@ -347,7 +353,7 @@ def __getitem__(self, key: Union[str, Enum]) -> Any:
         except Exception as e:
             self._format_and_raise(key=key, value=None, cause=e)
 
-    def __delitem__(self, key: Union[str, int, Enum]) -> None:
+    def __delitem__(self, key: DictKeyType) -> None:
         if self._get_flag("readonly"):
             self._format_and_raise(
                 key=key,
@@ -375,15 +381,13 @@ def __delitem__(self, key: Union[str, int, Enum]) -> None:
 
         del self.__dict__["_content"][key]
 
-    def get(
-        self, key: Union[str, Enum], default_value: Any = DEFAULT_VALUE_MARKER
-    ) -> Any:
+    def get(self, key: DictKeyType, default_value: Any = DEFAULT_VALUE_MARKER) -> Any:
         try:
             return self._get_impl(key=key, default_value=default_value)
         except Exception as e:
             self._format_and_raise(key=key, value=None, cause=e)
 
-    def _get_impl(self, key: Union[str, Enum], default_value: Any) -> Any:
+    def _get_impl(self, key: DictKeyType, default_value: Any) -> Any:
         try:
             node = self._get_node(key=key)
         except ConfigAttributeError:
@@ -396,7 +400,7 @@ def _get_impl(self, key: Union[str, Enum], default_value: Any) -> Any:
         )
 
     def _get_node(
-        self, key: Union[str, Enum], validate_access: bool = True
+        self, key: DictKeyType, validate_access: bool = True
     ) -> Optional[Node]:
         try:
             key = self._validate_and_normalize_key(key)
@@ -413,7 +417,7 @@ def _get_node(
 
         return value
 
-    def pop(self, key: Union[str, Enum], default: Any = DEFAULT_VALUE_MARKER) -> Any:
+    def pop(self, key: DictKeyType, default: Any = DEFAULT_VALUE_MARKER) -> Any:
         try:
             if self._get_flag("readonly"):
                 raise ReadonlyConfigError("Cannot pop from read-only node")
@@ -479,13 +483,13 @@ def __contains__(self, key: object) -> bool:
             except (MissingMandatoryValue, KeyError):
                 return False
 
-    def __iter__(self) -> Iterator[str]:
+    def __iter__(self) -> Iterator[DictKeyType]:
         return iter(self.keys())
 
-    def items(self) -> AbstractSet[Tuple[str, Any]]:
+    def items(self) -> AbstractSet[Tuple[DictKeyType, Any]]:
         return self.items_ex(resolve=True, keys=None)
 
-    def setdefault(self, key: Union[str, Enum], default: Any = None) -> Any:
+    def setdefault(self, key: DictKeyType, default: Any = None) -> Any:
         if key in self:
             ret = self.__getitem__(key)
         else:
@@ -494,9 +498,9 @@ def setdefault(self, key: Union[str, Enum], default: Any = None) -> Any:
         return ret
 
     def items_ex(
-        self, resolve: bool = True, keys: Optional[List[str]] = None
-    ) -> AbstractSet[Tuple[str, Any]]:
-        items: List[Tuple[str, Any]] = []
+        self, resolve: bool = True, keys: Optional[Sequence[DictKeyType]] = None
+    ) -> AbstractSet[Tuple[DictKeyType, Any]]:
+        items: List[Tuple[DictKeyType, Any]] = []
         for key in self.keys():
             if resolve:
                 value = self.get(key)

omegaconf/omegaconf.py

@@ -28,7 +28,7 @@
 import yaml
 from typing_extensions import Protocol
 
-from . import DictConfig, ListConfig
+from . import DictConfig, DictKeyType, ListConfig
 from ._utils import (
     _ensure_container,
     _get_value,
@@ -183,7 +183,7 @@ def create(
     @staticmethod
     @overload
     def create(
-        obj: Union[Dict[str, Any], None] = None,
+        obj: Optional[Dict[Any, Any]] = None,
         parent: Optional[BaseContainer] = None,
         flags: Optional[Dict[str, bool]] = None,
     ) -> DictConfig:
@@ -467,7 +467,7 @@ def to_container(
         resolve: bool = False,
         enum_to_str: bool = False,
         exclude_structured_configs: bool = False,
-    ) -> Union[Dict[str, Any], List[Any], None, str]:
+    ) -> Union[Dict[DictKeyType, Any], List[Any], None, str]:
         """
         Resursively converts an OmegaConf config to a primitive container (dict or list).
         :param cfg: the config to convert
@@ -508,7 +508,7 @@ def is_optional(obj: Any, key: Optional[Union[int, str]] = None) -> bool:
             return True
 
     @staticmethod
-    def is_none(obj: Any, key: Optional[Union[int, str]] = None) -> bool:
+    def is_none(obj: Any, key: Optional[Union[int, DictKeyType]] = None) -> bool:
         if key is not None:
             assert isinstance(obj, Container)
             obj = obj._get_node(key)

tests/structured_conf/data/attr_classes.py

@@ -300,9 +300,9 @@ class WithTypedDict:
 
 
 @attr.s(auto_attribs=True)
-class ErrorDictIntKey:
+class ErrorDictObjectKey:
     # invalid dict key, must be str
-    dict: Dict[int, str] = {10: "foo", 20: "bar"}
+    dict: Dict[object, str] = {object(): "foo", object(): "bar"}
 
 
 class RegularClass:
@@ -350,6 +350,7 @@ class DictExamples:
         "green": Color.GREEN,
         "blue": Color.BLUE,
     }
+    int_keys: Dict[int, str] = {1: "one", 2: "two"}
 
 
 @attr.s(auto_attribs=True)
@@ -372,6 +373,10 @@ class DictSubclass:
     class Str2Str(Dict[str, str]):
         pass
 
+    @attr.s(auto_attribs=True)
+    class Int2Str(Dict[int, str]):
+        pass
+
     @attr.s(auto_attribs=True)
     class Color2Str(Dict[Color, str]):
         pass