fixes #3 - colorized output for runner

awslimitchecker/runner.py

@@ -41,6 +41,7 @@
 import argparse
 import logging
 import json
+import termcolor
 
 from .version import _get_version, _get_project_url
 from .checker import AwsLimitChecker
@@ -49,6 +50,8 @@
 logging.basicConfig(level=logging.WARNING)
 logger = logging.getLogger()
 
+colorize = True
+
 
 def parse_args(argv):
     """
@@ -107,6 +110,8 @@ def parse_args(argv):
                    type=int, default=99,
                    help='default critical threshold (percentage of '
                    'limit); default: 99')
+    p.add_argument('--no-color', action='store_true', default=False,
+                   help='do not colorize output')
     p.add_argument('-v', '--verbose', dest='verbose', action='count',
                    default=0,
                    help='verbose output. specify twice for debug-level output.')
@@ -158,6 +163,13 @@ def show_usage(checker):
     print(dict2cols(data))
 
 
+def color_output(s, color):
+    global colorize
+    if not colorize:
+        return s
+    return termcolor.colored(s, color)
+
+
 def print_issue(service_name, limit, crits, warns):
     """
     :param service_name: the name of the service
@@ -171,13 +183,15 @@ def print_issue(service_name, limit, crits, warns):
     """
     usage_str = ''
     if len(crits) > 0:
-        usage_str += 'CRITICAL: '
-        usage_str += ', '.join([str(x) for x in sorted(crits)])
+        tmp = 'CRITICAL: '
+        tmp += ', '.join([str(x) for x in sorted(crits)])
+        usage_str += color_output(tmp, 'red')
     if len(warns) > 0:
         if len(crits) > 0:
             usage_str += ' '
-        usage_str += 'WARNING: '
-        usage_str += ', '.join([str(x) for x in sorted(warns)])
+        tmp = 'WARNING: '
+        tmp += ', '.join([str(x) for x in sorted(warns)])
+        usage_str += color_output(tmp, 'yellow')
     k = "{s}/{l}".format(
         s=service_name,
         l=limit.name,
@@ -225,6 +239,7 @@ def set_limit_overrides(checker, overrides):
 
 
 def console_entry_point():
+    global colorize
     args = parse_args(sys.argv[1:])
     if args.verbose == 1:
         logger.setLevel(logging.INFO)
@@ -236,6 +251,9 @@ def console_entry_point():
         logger.handlers[0].setFormatter(debug_formatter)
         logger.setLevel(logging.DEBUG)
 
+    if args.no_color:
+        colorize = False
+
     if args.version:
         print('awslimitchecker {v} (see <{s}> for source code)'.format(
             s=_get_project_url(),

awslimitchecker/tests/test_runner.py

@@ -43,6 +43,7 @@
 import logging
 import json
 from mock import patch, call, Mock, DEFAULT
+import termcolor
 
 import awslimitchecker.runner as runner
 import awslimitchecker.version as version
@@ -52,6 +53,14 @@
 from .support import sample_limits
 
 
+def red(s):
+    return termcolor.colored(s, 'red')
+
+
+def yellow(s):
+    return termcolor.colored(s, 'yellow')
+
+
 class TestAwsLimitCheckerRunner(object):
 
     def test_parse_args(self):
@@ -114,6 +123,9 @@ def test_parse_args_parser(self):
                                 type=int, default=99,
                                 help='default critical threshold (percentage '
                                 'of limit); default: 99'),
+            call().add_argument('--no-color', action='store_true',
+                                default=False,
+                                help='do not colorize output'),
             call().add_argument('-v', '--verbose', dest='verbose',
                                 action='count',
                                 default=0,
@@ -627,7 +639,7 @@ def test_print_issue_crit_one(self):
             [c1],
             []
         )
-        assert res == ('svcname/limitname', '(limit 12) CRITICAL: 56')
+        assert res == ('svcname/limitname', '(limit 12) ' + red('CRITICAL: 56'))
 
     def test_print_issue_crit_multi(self):
         mock_limit = Mock(spec_set=AwsLimit)
@@ -645,7 +657,7 @@ def test_print_issue_crit_multi(self):
             []
         )
         assert res == ('svcname/limitname',
-                       '(limit 5) CRITICAL: 8, 10, c2id=12')
+                       '(limit 5) ' + red('CRITICAL: 8, 10, c2id=12'))
 
     def test_print_issue_warn_one(self):
         mock_limit = Mock(spec_set=AwsLimit)
@@ -660,7 +672,8 @@ def test_print_issue_warn_one(self):
             [],
             [w1]
         )
-        assert res == ('svcname/limitname', '(limit 12) WARNING: 11')
+        assert res == ('svcname/limitname', '(limit 12) ' +
+                       yellow('WARNING: 11'))
 
     def test_print_issue_warn_multi(self):
         mock_limit = Mock(spec_set=AwsLimit)
@@ -678,8 +691,8 @@ def test_print_issue_warn_multi(self):
             [w1, w2, w3]
         )
         assert res == ('svcname/limitname',
-                       '(limit 12) WARNING: '
-                       'w2id=10, w3id=10, 11')
+                       '(limit 12) ' + yellow('WARNING: '
+                                              'w2id=10, w3id=10, 11'))
 
     def test_print_issue_both_one(self):
         mock_limit = Mock(spec_set=AwsLimit)
@@ -696,9 +709,9 @@ def test_print_issue_both_one(self):
             [w1]
         )
         assert res == ('svcname/limitname',
-                       '(limit 12) '
-                       'CRITICAL: 10 '
-                       'WARNING: w3id=10')
+                       '(limit 12) ' +
+                       red('CRITICAL: 10') + ' ' +
+                       yellow('WARNING: w3id=10'))
 
     def test_print_issue_both_multi(self):
         mock_limit = Mock(spec_set=AwsLimit)
@@ -719,6 +732,25 @@ def test_print_issue_both_multi(self):
             [w1, w2, w3]
         )
         assert res == ('svcname/limitname',
-                       '(limit 12) '
-                       'CRITICAL: 8, 10, c2id=12 '
-                       'WARNING: w2id=10, w3id=10, 11')
+                       '(limit 12) ' +
+                       red('CRITICAL: 8, 10, c2id=12') + ' ' +
+                       yellow('WARNING: w2id=10, w3id=10, 11'))
+
+    def test_entry_no_color(self):
+        argv = ['awslimitchecker', '--no-color']
+        with patch.object(sys, 'argv', argv):
+            with patch.multiple(
+                    'awslimitchecker.runner',
+                    AwsLimitChecker=DEFAULT,
+                    check_thresholds=DEFAULT,
+            ) as mocks:
+                mocks['check_thresholds'].return_value = 0
+                with pytest.raises(SystemExit) as excinfo:
+                    runner.console_entry_point()
+        assert excinfo.value.code == 0
+        assert runner.color_output('foo', 'red') == 'foo'
+        runner.colorize = True
+
+    def test_color_output(self):
+        assert runner.color_output('foo', 'yellow') == termcolor.colored(
+            'foo', 'yellow')

docs/build_generated_docs.py

@@ -185,8 +185,9 @@ def build_runner_examples():
             '--limit="EC2/EC2-VPC Elastic IPs"=200',
             '--list-defaults',
         ],
-        'check_thresholds': ['awslimitchecker'],
-        'check_thresholds_custom': ['awslimitchecker', '-W', '97', '--critical=98'],
+        'check_thresholds': ['awslimitchecker', '--no-color'],
+        'check_thresholds_custom': ['awslimitchecker', '-W', '97',
+                                    '--critical=98', '--no-color'],
         'iam_policy': ['awslimitchecker', '--iam-policy'],
     }
     results = {}

docs/source/cli_usage.rst

@@ -26,7 +26,7 @@ use as a Nagios-compatible plugin).
    (venv)$ awslimitchecker --help
    usage: awslimitchecker [-h] [-s] [-l] [--list-defaults] [-L LIMIT] [-u]
                           [--iam-policy] [-W WARNING_THRESHOLD]
-                          [-C CRITICAL_THRESHOLD] [-v] [-V]
+                          [-C CRITICAL_THRESHOLD] [--no-color] [-v] [-V]
    Report on AWS service limits and usage via boto, optionally warn about any
    services with usage nearing or exceeding their limits. For further help, see
    <http://awslimitchecker.readthedocs.org/>
@@ -53,6 +53,7 @@ use as a Nagios-compatible plugin).
      -C CRITICAL_THRESHOLD, --critical-threshold CRITICAL_THRESHOLD
                            default critical threshold (percentage of limit);
                            default: 99
+     --no-color            do not colorize output
      -v, --verbose         verbose output. specify twice for debug-level output.
      -V, --version         print version number and exit.
    awslimitchecker is AGPLv3-licensed Free Software. Anyone using this program,
@@ -189,7 +190,7 @@ threshold only, and another has crossed the critical threshold):
 
 .. code-block:: console
 
-   (venv)$ awslimitchecker
+   (venv)$ awslimitchecker --no-color
    AutoScaling/Auto Scaling Groups            (limit 20) CRITICAL: 38
    EC2/EC2-Classic Elastic IPs                (limit 5) CRITICAL: 5
    EC2/Running On-Demand EC2 instances        (limit 20) CRITICAL: 98
@@ -206,7 +207,7 @@ To set the warning threshold of 50% and a critical threshold of 75% when checkin
 
 .. code-block:: console
 
-   (venv)$ awslimitchecker -W 97 --critical=98
+   (venv)$ awslimitchecker -W 97 --critical=98 --no-color
    AutoScaling/Auto Scaling Groups            (limit 20) CRITICAL: 38
    EC2/EC2-Classic Elastic IPs                (limit 5) CRITICAL: 5
    EC2/Running On-Demand EC2 instances        (limit 20) CRITICAL: 98

setup.py

@@ -46,6 +46,7 @@
 
 requires = [
     'boto>=2.32.0',
+    'termcolor==1.1.0',
 ]
 
 classifiers = [