added confirmation function in logs.py

thefuck/entrypoints/fix_command.py

@@ -6,17 +6,17 @@ from .. import logs, types, const
 from ..conf import settings
 from ..corrector import get_corrected_commands
 from ..exceptions import EmptyCommand
-from ..ui import select_command
+from ..ui import select_command, confirm_command
 from ..utils import get_alias, get_all_executables
 
 
 def _get_raw_command(known_args):
     if known_args.force_command:
         return known_args.force_command
-    elif not os.environ.get('TF_HISTORY'):
+    elif not os.environ.get("TF_HISTORY"):
         return known_args.command
     else:
-        history = os.environ['TF_HISTORY'].split('\n')[::-1]
+        history = os.environ["TF_HISTORY"].split("\n")[::-1]
         alias = get_alias()
         executables = get_all_executables()
         for command in history:
@@ -29,20 +29,24 @@ def _get_raw_command(known_args):
 def fix_command(known_args):
     """Fixes previous command. Used when `thefuck` called without arguments."""
     settings.init(known_args)
-    with logs.debug_time('Total'):
-        logs.debug(u'Run with settings: {}'.format(pformat(settings)))
+    with logs.debug_time("Total"):
+        logs.debug(u"Run with settings: {}".format(pformat(settings)))
         raw_command = _get_raw_command(known_args)
 
         try:
             command = types.Command.from_raw_script(raw_command)
         except EmptyCommand:
-            logs.debug('Empty command, nothing to do')
+            logs.debug("Empty command, nothing to do")
             return
 
         corrected_commands = get_corrected_commands(command)
         selected_command = select_command(corrected_commands)
 
-        if selected_command:
+        confirmation = True
+        if selected_command.script == "reboot":
+            confirmation = confirm_command("Reboot System?")
+
+        if selected_command and confirmation:
             selected_command.run(command)
         else:
             sys.exit(1)

thefuck/logs.py

@@ -139,3 +139,14 @@ def version(thefuck_version, python_version, shell_info):
         u'The Fuck {} using Python {} and {}\n'.format(thefuck_version,
                                                        python_version,
                                                        shell_info))
+
+
+def confirmation(confirm):
+    if confirm is True:
+        sys.stderr.write(u"\n{bold}System Rebooting!{reset}".format(
+            bold=color(colorama.Style.BRIGHT),
+            reset=color(colorama.Style.RESET_ALL)))
+    else:
+        sys.stderr.write(u"\n{bold}Reboot Cancelled{reset}".format(
+            bold=color(colorama.Style.BRIGHT),
+            reset=color(colorama.Style.RESET_ALL)))

thefuck/ui.py

@@ -6,6 +6,7 @@ from .exceptions import NoRuleMatched
 from .system import get_key
 from .utils import get_alias
 from . import logs, const
+from .types import CorrectedCommand
 
 
 def read_actions():
@@ -14,13 +15,13 @@ def read_actions():
         key = get_key()
 
         # Handle arrows, j/k (qwerty), and n/e (colemak)
-        if key in (const.KEY_UP, const.KEY_CTRL_N, 'k', 'e'):
+        if key in (const.KEY_UP, const.KEY_CTRL_N, "k", "e"):
             yield const.ACTION_PREVIOUS
-        elif key in (const.KEY_DOWN, const.KEY_CTRL_P, 'j', 'n'):
+        elif key in (const.KEY_DOWN, const.KEY_CTRL_P, "j", "n"):
             yield const.ACTION_NEXT
-        elif key in (const.KEY_CTRL_C, 'q'):
+        elif key in (const.KEY_CTRL_C, "q"):
             yield const.ACTION_ABORT
-        elif key in ('\n', '\r'):
+        elif key in ("\n", "\r"):
             yield const.ACTION_SELECT
 
 
@@ -58,20 +59,16 @@ class CommandSelector(object):
 
 def select_command(corrected_commands):
     """Returns:
-
      - the first command when confirmation disabled;
      - None when ctrl+c pressed;
      - selected command.
-
     :type corrected_commands: Iterable[thefuck.types.CorrectedCommand]
     :rtype: thefuck.types.CorrectedCommand | None
-
     """
     try:
         selector = CommandSelector(corrected_commands)
     except NoRuleMatched:
-        logs.failed('No fucks given' if get_alias() == 'fuck'
-                    else 'Nothing found')
+        logs.failed("No fucks given" if get_alias() == "fuck" else "Nothing found")
         return
 
     if not settings.require_confirmation:
@@ -82,10 +79,10 @@ def select_command(corrected_commands):
 
     for action in read_actions():
         if action == const.ACTION_SELECT:
-            sys.stderr.write('\n')
+            sys.stderr.write("\n")
             return selector.value
         elif action == const.ACTION_ABORT:
-            logs.failed('\nAborted')
+            logs.failed("\nAborted")
             return
         elif action == const.ACTION_PREVIOUS:
             selector.previous()
@@ -93,3 +90,23 @@ def select_command(corrected_commands):
         elif action == const.ACTION_NEXT:
             selector.next()
             logs.confirm_text(selector.value)
+
+
+def confirm_command(confirmation_text):
+    """Returns:
+     - the first command when confirmation disabled;
+     - None when ctrl+c pressed;
+     - selected command.
+    :type corrected_commands: Iterable[thefuck.types.CorrectedCommand]
+    :rtype: thefuck.types.CorrectedCommand | None
+    """
+
+    logs.confirm_text(CorrectedCommand(confirmation_text, None, 0))
+
+    for action in read_actions():
+        if action == const.ACTION_SELECT:
+            logs.confirmation(True)
+            return True
+        elif action == const.ACTION_ABORT:
+            logs.confirmation(False)
+            return False