From 14963d67d17da648fdd61a9b438c56149d11c0cf Mon Sep 17 00:00:00 2001
From: snail-coupe <53581981+snail-coupe@users.noreply.github.com>
Date: Sun, 20 Nov 2022 14:17:27 +0000
Subject: [PATCH] Add option to pass in existing screen. Fixes #91

---
 README.md            |  1 +
 src/pick/__init__.py | 19 ++++++++++++++++---
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index a29c245..2dd7d38 100644
--- a/README.md
+++ b/README.md
@@ -56,6 +56,7 @@ interactive selection list in the terminal.
   multiple items by hitting SPACE
 - `min_selection_count`: (optional) for multi select feature to
   dictate a minimum of selected items before continuing
+- `screen`: (optional), if you are using `pick` within an existing curses application set this to your existing `screen` object.  It is assumed this has initialised in the standard way (e.g. via `curses.wrapper()`, or `curses.noecho(); curses.cbreak(); screen.kepad(True)`)
 
 ## Community Projects
 
diff --git a/src/pick/__init__.py b/src/pick/__init__.py
index 22bacc5..c54ce83 100644
--- a/src/pick/__init__.py
+++ b/src/pick/__init__.py
@@ -33,6 +33,7 @@ class Picker(Generic[OPTION_T]):
     min_selection_count: int = 0
     selected_indexes: List[int] = field(init=False, default_factory=list)
     index: int = field(init=False, default=0)
+    screen: Optional["curses._CursesWindow"] = None
 
     def __post_init__(self) -> None:
         if len(self.options) == 0:
@@ -110,7 +111,7 @@ def get_lines(self) -> Tuple[List, int]:
         current_line = self.index + len(title_lines) + 1
         return lines, current_line
 
-    def draw(self, screen) -> None:
+    def draw(self, screen: "curses._CursesWindow") -> None:
         """draw the curses ui on the screen, handle scroll if needed"""
         screen.clear()
 
@@ -133,7 +134,9 @@ def draw(self, screen) -> None:
 
         screen.refresh()
 
-    def run_loop(self, screen) -> Union[List[PICK_RETURN_T], PICK_RETURN_T]:
+    def run_loop(
+        self, screen: "curses._CursesWindow"
+    ) -> Union[List[PICK_RETURN_T], PICK_RETURN_T]:
         while True:
             self.draw(screen)
             c = screen.getch()
@@ -161,11 +164,19 @@ def config_curses(self) -> None:
             # Curses failed to initialize color support, eg. when TERM=vt100
             curses.initscr()
 
-    def _start(self, screen):
+    def _start(self, screen: "curses._CursesWindow"):
         self.config_curses()
         return self.run_loop(screen)
 
     def start(self):
+        if self.screen:
+            # Given an existing screen
+            # don't make any lasting changes
+            last_cur = curses.curs_set(0)
+            ret = self.run_loop(self.screen)
+            if last_cur:
+                curses.curs_set(last_cur)
+            return ret
         return curses.wrapper(self._start)
 
 
@@ -176,6 +187,7 @@ def pick(
     default_index: int = 0,
     multiselect: bool = False,
     min_selection_count: int = 0,
+    screen: Optional["curses._CursesWindow"] = None,
 ):
     picker: Picker = Picker(
         options,
@@ -184,5 +196,6 @@ def pick(
         default_index,
         multiselect,
         min_selection_count,
+        screen,
     )
     return picker.start()