Use colorama for colored output

README.md

@@ -193,7 +193,8 @@ The Fuck has a few settings parameters, they can be changed in `~/.thefuck/setti
 
 * `rules` – list of enabled rules, by default all;
 * `require_confirmation` – require confirmation before running new command, by default `False`; 
-* `wait_command` – max amount of time in seconds for getting previous command output.
+* `wait_command` – max amount of time in seconds for getting previous command output;
+* `no_colors` – disable colored output.
 
 ## Developing
 

setup.py

@@ -15,6 +15,6 @@ setup(name='thefuck',
                                       'tests', 'release']),
       include_package_data=True,
       zip_safe=False,
-      install_requires=['pathlib', 'psutil'],
+      install_requires=['pathlib', 'psutil', 'colorama'],
       entry_points={'console_scripts': [
           'thefuck = thefuck.main:main']})

tests/test_logs.py

@@ -0,0 +1,7 @@
+from mock import Mock
+from thefuck import logs
+
+
+def test_color():
+    assert logs.color('red', Mock(no_colors=False)) == 'red'
+    assert logs.color('red', Mock(no_colors=True)) == ''

tests/test_main.py

@@ -69,11 +69,11 @@ def test_get_matched_rule(capsys):
              main.Rule('', lambda *_: False, None),
              main.Rule('rule', Mock(side_effect=OSError('Denied')), None)]
     assert main.get_matched_rule(main.Command('ls', '', ''),
-                                 rules, None) is None
+                                 rules, Mock(no_colors=True)) is None
     assert main.get_matched_rule(main.Command('cd ..', '', ''),
-                                 rules, None) == rules[0]
+                                 rules, Mock(no_colors=True)) == rules[0]
     assert capsys.readouterr()[1].split('\n')[0]\
-           == '[WARN] rule: Traceback (most recent call last):'
+           == '[WARN] Rule rule:'
 
 
 def test_run_rule(capsys):
@@ -93,9 +93,11 @@ def test_confirm(capsys):
     assert capsys.readouterr() == ('', 'command\n')
     # When confirmation required and confirmed:
     with patch('thefuck.main.sys.stdin.read', return_value='\n'):
-        assert main.confirm('command', Mock(require_confirmation=True))
-        assert capsys.readouterr() == ('', 'command [Enter/Ctrl+C]')
+        assert main.confirm('command', Mock(require_confirmation=True,
+                                            no_colors=True))
+        assert capsys.readouterr() == ('', 'command [enter/ctrl+c]')
     # When confirmation required and ctrl+c:
     with patch('thefuck.main.sys.stdin.read', side_effect=KeyboardInterrupt):
-        assert not main.confirm('command', Mock(require_confirmation=True))
-        assert capsys.readouterr() == ('', 'command [Enter/Ctrl+C]Aborted\n')
+        assert not main.confirm('command', Mock(require_confirmation=True,
+                                                no_colors=True))
+        assert capsys.readouterr() == ('', 'command [enter/ctrl+c]Aborted\n')

thefuck/logs.py

@@ -0,0 +1,47 @@
+import sys
+from traceback import format_exception
+import colorama
+
+
+def color(color_, settings):
+    """Utility for ability to disabling colored output."""
+    if settings.no_colors:
+        return ''
+    else:
+        return color_
+
+
+def rule_failed(rule, exc_info, settings):
+    sys.stderr.write(
+        u'{warn}[WARN] Rule {name}:{reset}\n{trace}'
+        u'{warn}----------------------------{reset}\n\n'.format(
+            warn=color(colorama.Back.RED + colorama.Fore.WHITE
+                       + colorama.Style.BRIGHT, settings),
+            reset=color(colorama.Style.RESET_ALL, settings),
+            name=rule.name,
+            trace=''.join(format_exception(*exc_info))))
+
+
+def show_command(new_command, settings):
+    sys.stderr.write('{bold}{command}{reset}\n'.format(
+        command=new_command,
+        bold=color(colorama.Style.BRIGHT, settings),
+        reset=color(colorama.Style.RESET_ALL, settings)))
+
+
+def confirm_command(new_command, settings):
+    sys.stderr.write(
+        '{bold}{command}{reset} [{green}enter{reset}/{red}ctrl+c{reset}]'.format(
+            command=new_command,
+            bold=color(colorama.Style.BRIGHT, settings),
+            green=color(colorama.Fore.GREEN, settings),
+            red=color(colorama.Fore.RED, settings),
+            reset=color(colorama.Style.RESET_ALL, settings)))
+    sys.stderr.flush()
+
+
+def failed(msg, settings):
+    sys.stderr.write('{red}{msg}{reset}\n'.format(
+        msg=msg,
+        red=color(colorama.Fore.RED, settings),
+        reset=color(colorama.Style.RESET_ALL, settings)))

thefuck/main.py

@@ -5,8 +5,9 @@ from os.path import expanduser
 from subprocess import Popen, PIPE
 import os
 import sys
-from traceback import format_exception
 from psutil import Process, TimeoutExpired
+import colorama
+from thefuck import logs
 
 
 Command = namedtuple('Command', ('script', 'stdout', 'stderr'))
@@ -30,6 +31,7 @@ def get_settings(user_dir):
     settings.__dict__.setdefault('rules', None)
     settings.__dict__.setdefault('wait_command', 3)
     settings.__dict__.setdefault('require_confirmation', False)
+    settings.__dict__.setdefault('no_colors', False)
     return settings
 
 
@@ -100,23 +102,21 @@ def get_matched_rule(command, rules, settings):
             if rule.match(command, settings):
                 return rule
         except Exception:
-            sys.stderr.write(u'[WARN] {}: {}---------------------\n\n'.format(
-                rule.name, ''.join(format_exception(*sys.exc_info()))))
+            logs.rule_failed(rule, sys.exc_info(), settings)
 
 
 def confirm(new_command, settings):
     """Returns `True` when running of new command confirmed."""
     if not settings.require_confirmation:
-        sys.stderr.write(new_command + '\n')
+        logs.show_command(new_command, settings)
         return True
 
-    sys.stderr.write(new_command + ' [Enter/Ctrl+C]')
-    sys.stderr.flush()
+    logs.confirm_command(new_command, settings)
     try:
         sys.stdin.read(1)
         return True
     except KeyboardInterrupt:
-        sys.stderr.write('Aborted\n')
+        logs.failed('Aborted', settings)
         return False
 
 
@@ -133,13 +133,14 @@ def is_second_run(command):
 
 
 def main():
+    colorama.init()
     user_dir = setup_user_dir()
     settings = get_settings(user_dir)
 
     command = get_command(settings, sys.argv)
     if command:
         if is_second_run(command):
-            sys.stderr.write("Can't fuck twice\n")
+            logs.failed("Can't fuck twice", settings)
             return
 
         rules = get_rules(user_dir, settings)
@@ -148,4 +149,4 @@ def main():
             run_rule(matched_rule, command, settings)
             return
 
-    sys.stderr.write('No fuck given\n')
+    logs.failed('No fuck given', settings)