From 68b645b038c4eb99205c39c13aff1712804d67fd Mon Sep 17 00:00:00 2001
From: Trevor Taubitz <trevor.taubitz@gmail.com>
Date: Thu, 28 Dec 2017 14:25:27 -0500
Subject: [PATCH] Added boolean prompt type

---
 prompt/__init__.py     | 41 +++++++++++++++++++++++++++++++++++++++++
 tests/test_expected.py | 11 +++++++++++
 2 files changed, 52 insertions(+)

diff --git a/prompt/__init__.py b/prompt/__init__.py
index e5bd945..916a018 100644
--- a/prompt/__init__.py
+++ b/prompt/__init__.py
@@ -260,6 +260,47 @@ def string(prompt=None, empty=False):
             return string(prompt=prompt, empty=empty)
 
 
+def boolean(prompt=None, default=None):
+    """Prompt a yes or no question.
+
+    Parameters
+    ----------
+    prompt : str, optional
+        Use an alternative prompt.
+    default : bool, optional
+        Set a default response
+
+    Returns
+    -------
+    bool
+        A bool indicating the response
+    """
+
+    if not prompt:
+        prompt = PROMPT
+    if default is None:
+        prompt_addition = ' (y/n): '
+    elif default:
+        prompt_addition = ' (Y/n): '
+    else:
+        prompt_addition = ' (y/N): '
+
+    prompt += prompt_addition
+    response = None
+    while response is None:
+        # Keep asking for a response until either the default or a proper char is used
+        char = character(prompt, empty=True)
+        if char is None:
+            # Use the default if no response. If no default is set, the loop will repeat
+            response = default
+        elif char.lower() == 'y':
+            response = True
+        elif char.lower() == 'n':
+            response = False
+
+    return response
+
+
 def _prompt_input(prompt):
     if prompt is None:
         return input(PROMPT)
diff --git a/tests/test_expected.py b/tests/test_expected.py
index 43820dc..236db97 100644
--- a/tests/test_expected.py
+++ b/tests/test_expected.py
@@ -62,3 +62,14 @@ def test_secret(input_patch):
 def test_string(input_patch):
     input_patch.do("foo123")
     assert prompt.string() == "foo123"
+
+
+def test_boolean(input_patch):
+    input_patch.do("Y")
+    assert prompt.boolean() == True
+    input_patch.do("y")
+    assert prompt.boolean() == True
+    input_patch.do("N")
+    assert prompt.boolean() == False
+    input_patch.do("n")
+    assert prompt.boolean() == False