From 8b465623d127e63878cc392d405b979a75ca14e1 Mon Sep 17 00:00:00 2001
From: skullY <skullydazed@gmail.com>
Date: Mon, 1 Jul 2019 01:31:29 -0700
Subject: [PATCH] fix qmk setup when run on a fresh system

---
 qmk_cli/helpers.py           |  6 +++---
 qmk_cli/script_qmk.py        |  2 +-
 qmk_cli/subcommands/setup.py | 18 +++++++++++++-----
 3 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/qmk_cli/helpers.py b/qmk_cli/helpers.py
index 4fea864..237c2c0 100644
--- a/qmk_cli/helpers.py
+++ b/qmk_cli/helpers.py
@@ -3,12 +3,12 @@
 from qmk_cli.milc import cli
 
 
-def question(question, boolean=True, default=None):
+def question(question, boolean=True, default=''):
     """Asks the user to answer a question.
 
     This keeps re-asking until it gets acceptible input.
     """
-    if cli.args.yes:
+    if cli.args.general_yes:
         return True
 
     if default and default.lower() == 'y':
@@ -18,7 +18,7 @@ def question(question, boolean=True, default=None):
     else:
         answer_key = 'y/n'
 
-    prompt = '*** %s [%s] ' % (answer_key,)
+    prompt = '*** %s [%s] ' % (question, answer_key)
 
     while True:
         answer = input(prompt)
diff --git a/qmk_cli/script_qmk.py b/qmk_cli/script_qmk.py
index f05f61b..07247ff 100644
--- a/qmk_cli/script_qmk.py
+++ b/qmk_cli/script_qmk.py
@@ -82,7 +82,7 @@ def main():
         import_module(subcommand_module)
         qmk_cli.milc.cli()
 
-    except ModuleNotFoundError as e:
+    except ImportError as e:
         # Check to make sure there's not a bad import statement in qmk_cli
         if e.name != subcommand_module:
             raise
diff --git a/qmk_cli/subcommands/setup.py b/qmk_cli/subcommands/setup.py
index 5e0b124..e9adf78 100644
--- a/qmk_cli/subcommands/setup.py
+++ b/qmk_cli/subcommands/setup.py
@@ -21,24 +21,32 @@ default_branch = 'master'
 @cli.argument('fork', default=default_fork, nargs='?', help='The qmk_firmware fork to clone')
 @cli.entrypoint('Setup your computer for qmk_firmware.')
 def main(cli):
+    setup_successful = False
     qmk_firmware = Path(cli.args.destination)
 
     # Check on qmk_firmware, and if it doesn't exist offer to check it out.
     if qmk_firmware.exists():
         cli.log.info('Found qmk_firmware at %s.', str(qmk_firmware))
     else:
+        cli.log.error('qmk_firmware not found!')
         if question('Would you like to clone %s?' % cli.args.fork):
             git_url = '/'.join((cli.config.general.baseurl, cli.args.fork))
             clone(git_url, cli.args.destination, cli.config.general.branch)
 
     # Check if the build environment is setup, and if not offer to set it up
-    if not check_vital_programs():
-        cli.log.info('Your build environment is not setup completely.')
-        if question('Would you like to run util/qmk_install?'):
+    if check_vital_programs():
+        cli.log.info('Your build environment is ready!')
+    else:
+        cli.log.error('Your build environment is not setup completely.')
+        if qmk_firmware.exists() and question('Would you like to run util/qmk_install?'):
             curdir = os.getcwd()
             os.chdir(str(qmk_firmware))
-            subprocess.run(['util/qmk_install.sh'])
+            process = subprocess.run(['util/qmk_install.sh'])
             os.chdir(curdir)
+            if process.returncode == 0:
+                setup_successful = True
+
 
     # fin
-    cli.log.info('QMK setup complete!')
+    if setup_successful:
+        cli.log.info('QMK setup complete!')