Update __init__.py

Lexanius · Lexanius · commit bd861933f476 · 2024-01-20T23:45:00.000Z
- Add check to select single values from choices
- Add test for single value selection
diff --git a/cli_ui/__init__.py b/cli_ui/__init__.py
@@ -591,7 +591,11 @@ def select_choices(
             continue
 
         try:
-            res = list(itemgetter(*index)(choices))
+            res = (
+                list(itemgetter(*index)(choices))
+                if len(index) > 1
+                else [itemgetter(*index)(choices)]
+            )
         except Exception:
             info("Please enter valid selection number(s)")
             continue
diff --git a/cli_ui/tests/test_cli_ui.py b/cli_ui/tests/test_cli_ui.py
@@ -345,6 +345,13 @@ def func_desc(fruit: Fruit) -> str:
         assert m.call_count == 3
 
 
+def test_select_choices_single_select() -> None:
+    with mock.patch("builtins.input") as m:
+        m.side_effect = ["1"]
+        res = cli_ui.select_choices("Select a animal", choices=["cat", "dog", "cow"])
+        assert res == ["cat"]
+
+
 def test_select_choices_empty_input() -> None:
     with mock.patch("builtins.input") as m:
         m.side_effect = [""]
