Skip to content

Commit

Permalink
dependency: replace colorama with custom module
Browse files Browse the repository at this point in the history
* In stretch there is a regression in colorama in conjunction with
  tqdm that leads to a slow down of the progress of the script
  proportional to the amount of data printed to stdout/err. Colorama
  starts having very huge stacktraces and the process is stuck at 100%
  CPU for an increasingly amount of time while more data is printed.
* Given the very simple usage of colors that is made in Cumin as of
  now, it seems much more feasible to replace the colorama library
  (as all that cross-OS support is not needed) and add a simple module
  with ANSI escape sequence support.
* Use a type (metaclass) to be able to override __getattr__ for the
  staticmethods of the classes that use it and to automatically define a
  method for each color in a DRY way without code duplication.
* Define a Colored class that uses ColoredType as metaclass to inherit
  its type with the custom behaviour.
* For each color defined in ColoredType.COLORS a method of Colored is
  defined, e.g. Colored.red().
* The Colored class has a `disabled` property that can be set to True to
  globally disable coloring. This could for example be integrated later
  into the CLI as an option to disable colors or allow to add some code
  to the color.py module to autodetect when not in a TTY and
  automatically disable all colors.

Bug: T217038
Change-Id: I053972e92b7cc6cf7821de8ecff484d01be03794
  • Loading branch information
volans- committed Jul 19, 2019
1 parent 79d8301 commit efb86bc
Show file tree
Hide file tree
Showing 7 changed files with 145 additions and 47 deletions.
9 changes: 3 additions & 6 deletions cumin/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,13 +11,12 @@

from logging.handlers import RotatingFileHandler # pylint: disable=ungrouped-imports

import colorama

from tqdm import tqdm

import cumin

from cumin import backends, query, transport, transports
from cumin.color import Colored


logger = logging.getLogger(__name__) # pylint: disable=invalid-name
Expand Down Expand Up @@ -292,8 +291,7 @@ def stderr(message, end='\n'):
end: the character to use at the end of the message. [optional, default: \n]
"""
tqdm.write('{color}{message}{reset}'.format(
color=colorama.Fore.YELLOW, message=message, reset=colorama.Style.RESET_ALL), file=sys.stderr, end=end)
tqdm.write(Colored.yellow(message), file=sys.stderr, end=end)


def get_hosts(args, config):
Expand All @@ -311,7 +309,7 @@ def get_hosts(args, config):
return hosts

stderr('{num} hosts will be targeted:'.format(num=len(hosts)))
stderr('{color}{hosts}'.format(color=colorama.Fore.CYAN, hosts=cumin.nodeset_fromlist(hosts)))
stderr(Colored.cyan(cumin.nodeset_fromlist(hosts)))

if args.dry_run:
stderr('DRY-RUN mode enabled, aborting')
Expand Down Expand Up @@ -439,7 +437,6 @@ def main(argv=None):
argv = sys.argv[1:]

signal.signal(signal.SIGINT, sigint_handler)
colorama.init()

# Setup
try:
Expand Down
61 changes: 61 additions & 0 deletions cumin/color.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
"""Colors module."""
from abc import ABCMeta


class ColoredType(ABCMeta):
"""Metaclass to define a new type that dynamically adds static methods to its classes."""

COLORS = {
'red': 31,
'green': 32,
'yellow': 33,
'blue': 34,
'cyan': 36,
}
""":py:class:`dict`: a mapping of colors to the ANSI foreground color code."""

def __getattr__(cls, name): # noqa: N805 (Prospector reqires an older version of pep8-naming)
"""Dynamically check access to members of classes of this type.
:Parameters:
according to Python's Data model :py:meth:`object.__getattr__`.
"""
color_code = ColoredType.COLORS.get(name, None)
if color_code is None:
raise AttributeError("'{cls}' object has no attribute '{attr}'".format(cls=cls.__name__, attr=name))

return lambda obj: cls._color(color_code, obj)


class Colored(metaclass=ColoredType):
"""Class to manage colored output.
Available methods are dynamically added based on the keys of the :py:const:`ColoredType.COLORS` dictionary.
For each color a method with the color name is available to color any object with that specific color code.
"""

disabled = False
""":py:class:`bool`: switch to globally control the coloring. Set it to :py:const`True` to disable all coloring."""

@staticmethod
def _color(color_code, obj):
"""Color the given object, unless coloring is globally disabled.
Arguments:
color_code (int): a valid ANSI escape sequence color code.
obj (mixed): the object to color.
Return:
str: the string representation of the object encapsulated in the red ANSI escape sequence.
"""
message = str(obj)

if not message:
return ''

if Colored.disabled:
return message

return '\x1b[{code}m{message}\x1b[39m'.format(code=color_code, message=message)
56 changes: 56 additions & 0 deletions cumin/tests/unit/test_color.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
"""Color tests."""
from unittest import mock

import pytest

from cumin.color import Colored


def test_red():
"""It should return the message enclosed in ASCII red color code."""
assert Colored.red('message') == '\x1b[31mmessage\x1b[39m'


def test_green():
"""It should return the message enclosed in ASCII green color code."""
assert Colored.green('message') == '\x1b[32mmessage\x1b[39m'


def test_yellow():
"""It should return the message enclosed in ASCII yellow color code."""
assert Colored.yellow('message') == '\x1b[33mmessage\x1b[39m'


def test_blue():
"""It should return the message enclosed in ASCII blue color code."""
assert Colored.blue('message') == '\x1b[34mmessage\x1b[39m'


def test_cyan():
"""It should return the message enclosed in ASCII cyan color code."""
assert Colored.cyan('message') == '\x1b[36mmessage\x1b[39m'


def test_wrong_case():
"""It should raise AttributeError if called with the wrong case."""
with pytest.raises(AttributeError, match="'Colored' object has no attribute 'Red'"):
Colored.Red('')


def test_non_existent():
"""It should raise AttributeError if called with a non existent color."""
with pytest.raises(AttributeError, match="'Colored' object has no attribute 'missing'"):
Colored.missing('')


def test_emtpy():
"""It should return an empty string if the object is empty."""
assert Colored.red('') == ''


@mock.patch('cumin.color.Colored.disabled', new_callable=mock.PropertyMock)
def test_disabled(mocked_colored_disabled):
"""It should return the message untouched if coloration is disabled."""
mocked_colored_disabled.return_value = True
assert Colored.red('message') == 'message'
assert mocked_colored_disabled.called
19 changes: 5 additions & 14 deletions cumin/tests/unit/transports/test_clustershell.py
Original file line number Diff line number Diff line change
Expand Up @@ -191,13 +191,11 @@ def setup_method(self, *args): # pylint: disable=arguments-differ
self.handler = None
self.args = args

@mock.patch('cumin.transports.clustershell.colorama')
def test_close(self, colorama):
def test_close(self):
"""Calling close should raise NotImplementedError."""
self.handler = clustershell.BaseEventHandler(self.target, self.commands)
with pytest.raises(NotImplementedError):
self.handler.close(self.worker)
colorama.init.assert_called_once_with(autoreset=True)


class ConcreteBaseEventHandler(clustershell.BaseEventHandler):
Expand All @@ -215,20 +213,17 @@ def close(self, task):
class TestConcreteBaseEventHandler(TestBaseEventHandler):
"""ConcreteBaseEventHandler test class."""

@mock.patch('cumin.transports.clustershell.colorama')
@mock.patch('cumin.transports.clustershell.tqdm')
def setup_method(self, _, tqdm, colorama): # pylint: disable=arguments-differ
def setup_method(self, _, tqdm): # pylint: disable=arguments-differ
"""Initialize default properties and instances."""
super().setup_method()
self.handler = ConcreteBaseEventHandler(self.target, self.commands)
self.worker.eh = self.handler
self.colorama = colorama
assert not tqdm.write.called

def test_instantiation(self):
"""An instance of ConcreteBaseEventHandler should be an instance of BaseEventHandler and initialize colorama."""
"""An instance of ConcreteBaseEventHandler should be an instance of BaseEventHandler."""
assert sorted(self.handler.nodes.keys()) == list(self.target.hosts)
self.colorama.init.assert_called_once_with(autoreset=True)

@mock.patch('cumin.transports.clustershell.tqdm')
def test_on_timeout(self, tqdm):
Expand Down Expand Up @@ -286,15 +281,13 @@ class TestSyncEventHandler(TestBaseEventHandler):
"""SyncEventHandler test class."""

@mock.patch('cumin.transports.clustershell.logging')
@mock.patch('cumin.transports.clustershell.colorama')
@mock.patch('cumin.transports.clustershell.tqdm')
def setup_method(self, _, tqdm, colorama, logger): # pylint: disable=arguments-differ
def setup_method(self, _, tqdm, logger): # pylint: disable=arguments-differ
"""Initialize default properties and instances."""
super().setup_method()
self.handler = clustershell.SyncEventHandler(self.target, self.commands, success_threshold=1)
self.handler.progress = mock.Mock(spec_set=ProgressBars)
self.worker.eh = self.handler
self.colorama = colorama
self.logger = logger
assert not tqdm.write.called

Expand Down Expand Up @@ -381,14 +374,12 @@ class TestAsyncEventHandler(TestBaseEventHandler):

@mock.patch('cumin.transports.clustershell.ProgressBars')
@mock.patch('cumin.transports.clustershell.logging')
@mock.patch('cumin.transports.clustershell.colorama')
@mock.patch('cumin.transports.clustershell.tqdm')
def setup_method(self, _, tqdm, colorama, logger, progress): # pylint: disable=arguments-differ,unused-argument
def setup_method(self, _, tqdm, logger, progress): # pylint: disable=arguments-differ,unused-argument
"""Initialize default properties and instances."""
super().setup_method()
self.handler = clustershell.AsyncEventHandler(self.target, self.commands)
self.worker.eh = self.handler
self.colorama = colorama
self.logger = logger
assert not tqdm.write.called

Expand Down
7 changes: 3 additions & 4 deletions cumin/transports/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,12 +6,11 @@

from abc import ABCMeta, abstractmethod, abstractproperty

import colorama

from ClusterShell.NodeSet import NodeSet
from tqdm import tqdm

from cumin import CuminError, nodeset_fromlist
from cumin.color import Colored


class WorkerError(CuminError):
Expand Down Expand Up @@ -751,10 +750,10 @@ def init(self, num_hosts):
"""
self.pbar_ok = tqdm(desc='PASS', total=num_hosts, leave=True, unit='hosts', dynamic_ncols=True,
bar_format=colorama.Fore.GREEN + self.bar_format, file=sys.stderr)
bar_format=Colored.green(self.bar_format), file=sys.stderr)
self.pbar_ok.refresh()
self.pbar_ko = tqdm(desc='FAIL', total=num_hosts, leave=True, unit='hosts', dynamic_ncols=True,
bar_format=colorama.Fore.RED + self.bar_format, file=sys.stderr)
bar_format=Colored.red(self.bar_format), file=sys.stderr)
self.pbar_ko.refresh()

def close(self):
Expand Down
39 changes: 17 additions & 22 deletions cumin/transports/clustershell.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,12 +5,11 @@

from collections import Counter, defaultdict

import colorama

from ClusterShell import Event, Task
from tqdm import tqdm

from cumin import nodeset, nodeset_fromlist
from cumin.color import Colored
from cumin.transports import BaseWorker, NoProgress, ProgressBars, raise_error, State, WorkerError


Expand Down Expand Up @@ -190,8 +189,7 @@ def __init__(self, target, commands, success_threshold=1.0, progress_bars=True,
for node_name in target.first_batch:
self.nodes[node_name].state.update(State.scheduled)

# Initialize color and progress bar formats
colorama.init(autoreset=True)
# Initialize progress bar formats
if progress_bars:
self.progress = ProgressBars()
else:
Expand Down Expand Up @@ -257,7 +255,7 @@ def ev_pickup(self, worker, node):
if not self.deduplicate_output:
output_message = "----- OUTPUT of '{command}' -----".format(
command=self._get_short_command(worker.command))
tqdm.write(colorama.Fore.BLUE + output_message, file=sys.stdout)
tqdm.write(Colored.blue(output_message), file=sys.stdout)

def ev_read(self, worker, node, _, msg):
"""Worker has data to read from a specific node. Print it if running on a single host.
Expand Down Expand Up @@ -322,18 +320,16 @@ def _get_log_message(self, num, message, nodes=None):

return (log_message, str(nodes_string))

def _print_report_line(self, message, color=colorama.Fore.RED, nodes_string=''): # pylint: disable=no-self-use
def _print_report_line(self, message, color_func=Colored.red, nodes_string=''): # pylint: disable=no-self-use
"""Print a tqdm-friendly colored status line with success/failure ratio and optional list of nodes.
Arguments:
message (str): the message to print.
color (str, optional): the message color.
color_func (function, optional): the coloring function, one of :py:class`cumin.color.Colored` methods.
nodes_string (str, optional): the string representation of the affected nodes.
"""
tqdm.write('{color}{message}{nodes_color}{nodes_string}'.format(
color=color, message=message, nodes_color=colorama.Fore.CYAN,
nodes_string=nodes_string), file=sys.stderr)
tqdm.write(color_func(message) + Colored.cyan(nodes_string), file=sys.stderr)

def _get_short_command(self, command):
"""Return a shortened representation of a command omitting the central part, if it's too long.
Expand All @@ -358,7 +354,7 @@ def _commands_output_report(self, buffer_iterator, command=None):
"""
if not self.deduplicate_output:
tqdm.write(colorama.Fore.BLUE + '================', file=sys.stdout)
tqdm.write(Colored.blue('================'), file=sys.stdout)
return

nodelist = None
Expand All @@ -368,19 +364,18 @@ def _commands_output_report(self, buffer_iterator, command=None):
output_message = '----- OUTPUT -----'

for output, nodelist in buffer_iterator.iter_buffers():
tqdm.write(colorama.Fore.BLUE + '===== NODE GROUP =====', file=sys.stdout)
tqdm.write('{color}({num}) {nodes}'.format(
color=colorama.Fore.CYAN, num=len(nodelist),
nodes=nodeset_fromlist(nodelist)), file=sys.stdout)
tqdm.write(colorama.Fore.BLUE + output_message, file=sys.stdout)
tqdm.write('{output}'.format(output=output.message().decode()), file=sys.stdout)
tqdm.write(Colored.blue('===== NODE GROUP ====='), file=sys.stdout)
tqdm.write(Colored.cyan('({num}) {nodes}'.format(num=len(nodelist), nodes=nodeset_fromlist(nodelist))),
file=sys.stdout)
tqdm.write(Colored.blue(output_message), file=sys.stdout)
tqdm.write(output.message().decode(), file=sys.stdout)

if nodelist is None:
message = '===== NO OUTPUT ====='
else:
message = '================'

tqdm.write(colorama.Fore.BLUE + message, file=sys.stdout)
tqdm.write(Colored.blue(message), file=sys.stdout)

def _global_timeout_nodes_report(self):
"""Print the nodes that were caught by the global timeout in a colored and tqdm-friendly way."""
Expand Down Expand Up @@ -460,17 +455,17 @@ def _success_nodes_report(self, command=None):
log_message, nodes_string = self._get_log_message(num, message, nodes=nodes)

if num == tot:
color = colorama.Fore.GREEN
color_func = Colored.green
level = logging.INFO
elif success_ratio >= self.success_threshold:
color = colorama.Fore.YELLOW
color_func = Colored.yellow
level = logging.WARNING
else:
color = colorama.Fore.RED
color_func = Colored.red
level = logging.CRITICAL

self.logger.log(level, '%s%s', log_message, nodes_string)
self._print_report_line(log_message, color=color, nodes_string=nodes_string)
self._print_report_line(log_message, color_func=color_func, nodes_string=nodes_string)


class SyncEventHandler(BaseEventHandler):
Expand Down
1 change: 0 additions & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@
# Required dependencies
install_requires = [
'clustershell==1.8',
'colorama>=0.3.2',
'pyparsing==2.1.10',
'pyyaml>=3.11',
'requests>=2.11.0',
Expand Down

0 comments on commit efb86bc

Please sign in to comment.