[Enhance] Add support for dict input in Choice

siatune/tune/spaces/choice.py

@@ -1,5 +1,5 @@
 # Copyright (c) SI-Analytics. All rights reserved.
-from typing import Callable, Optional, Sequence
+from typing import Callable, Optional, Sequence, Union
 
 import ray.tune as tune
 from ray.tune.search.sample import Domain
@@ -14,16 +14,19 @@ class Choice(BaseSpace):
     """Sample a categorical value.
 
     Args:
-        categories (Sequence): The categories.
-        alias (Sequence, optional): A alias to be expressed.
-            Defaults to None.
+        categories (Sequence | dict): The categorical search space to choose
+            one. If categories is dict, keys of dict will overwrite the alias.
+        alias (Sequence, optional): A alias to be expressed. Defaults to None.
     """
 
     sample: Callable = tune.choice
 
     def __init__(self,
-                 categories: Sequence,
+                 categories: Union[Sequence, dict],
                  alias: Optional[Sequence] = None) -> None:
+        if isinstance(categories, dict):
+            alias, categories = zip(*[(k, v) for k, v in categories.items()])
+
         if alias is not None:
             assert isinstance(alias, Sequence)
             assert len(categories) == len(alias)

tests/test_tune/test_spaces.py

@@ -33,6 +33,7 @@ def is_immutable(config):
     with pytest.raises(AssertionError):
         choice = Choice(categories=[True, False], alias=['TF'])
     choice = Choice(categories=[True, False], alias=['T', 'F'])
+    choice = Choice(categories=dict(T=True, F=False))
     tune.run(is_immutable, config=dict(test=choice.space))