From a043d6699f40818c0482e09a1f2fe3952d2e8c40 Mon Sep 17 00:00:00 2001
From: wong2 <wonderfuly@gmail.com>
Date: Wed, 31 Aug 2022 15:31:02 +0800
Subject: [PATCH] Add Option class

---
 example/option.py    | 10 ++++++++++
 src/pick/__init__.py | 22 ++++++++++++++--------
 tests/test_pick.py   | 12 +++++++++++-
 3 files changed, 35 insertions(+), 9 deletions(-)
 create mode 100644 example/option.py

diff --git a/example/option.py b/example/option.py
new file mode 100644
index 0000000..27b882c
--- /dev/null
+++ b/example/option.py
@@ -0,0 +1,10 @@
+from pick import pick, Option
+
+title = "Please choose your favorite programming language: "
+options = [
+    Option("Java", ".java"),
+    Option("Python", ".py"),
+    Option("JavaScript", ".js"),
+]
+option, index = pick(options, title)
+print(f"You choosed {option} at index {index}")
diff --git a/src/pick/__init__.py b/src/pick/__init__.py
index 710ad32..859fc38 100644
--- a/src/pick/__init__.py
+++ b/src/pick/__init__.py
@@ -1,8 +1,14 @@
 import curses
 from dataclasses import dataclass, field
-from typing import List, Optional, Sequence, Tuple, Union
+from typing import Any, List, Optional, Sequence, Tuple, TypeVar, Union, Generic
 
-__all__ = ["Picker", "pick"]
+__all__ = ["Picker", "pick", "Option"]
+
+
+@dataclass
+class Option:
+    label: str
+    value: Any
 
 
 KEYS_ENTER = (curses.KEY_ENTER, ord("\n"), ord("\r"))
@@ -13,13 +19,13 @@
 SYMBOL_CIRCLE_FILLED = "◉"
 SYMBOL_CIRCLE_EMPTY = "◯"
 
-KEY_T = int
-PICK_RETURN_T = Tuple[str, int]
+OPTION_T = TypeVar("OPTION_T", str, Option)
+PICK_RETURN_T = Tuple[OPTION_T, int]
 
 
 @dataclass
-class Picker:
-    options: Sequence[str]
+class Picker(Generic[OPTION_T]):
+    options: Sequence[OPTION_T]
     title: Optional[str] = None
     indicator: str = "*"
     default_index: int = 0
@@ -93,7 +99,7 @@ def get_option_lines(self) -> List[str]:
                 )
                 prefix = f"{prefix} {symbol} "
 
-            option_as_str = str(option)
+            option_as_str = option.label if isinstance(option, Option) else option
             lines.append(f"{prefix} {option_as_str}")
 
         return lines
@@ -166,7 +172,7 @@ def start(self):
 
 
 def pick(
-    options: Sequence[str],
+    options: Sequence[OPTION_T],
     title: Optional[str] = None,
     indicator: str = "*",
     default_index: int = 0,
diff --git a/tests/test_pick.py b/tests/test_pick.py
index 41401aa..c7a2e11 100644
--- a/tests/test_pick.py
+++ b/tests/test_pick.py
@@ -1,4 +1,4 @@
-from pick import Picker
+from pick import Picker, Option
 
 
 def test_move_up_down():
@@ -45,3 +45,13 @@ def test_multi_select():
     picker.move_down()
     picker.mark_index()
     assert picker.get_selected() == [("option1", 0), ("option2", 1)]
+
+
+def test_option():
+    options = [Option("option1", 101), Option("option2", 102), Option("option3", 103)]
+    picker = Picker(options)
+    picker.move_down()
+    option, index = picker.get_selected()
+    assert index == 1
+    assert isinstance(option, Option)
+    assert option.value == 102