diff --git a/colcon_core/output_style/__init__.py b/colcon_core/output_style/__init__.py index 9efc6fdb..c658e4d8 100644 --- a/colcon_core/output_style/__init__.py +++ b/colcon_core/output_style/__init__.py @@ -3,6 +3,7 @@ from collections import namedtuple import os +import re from types import SimpleNamespace from colcon_core.environment_variable import EnvironmentVariable @@ -14,6 +15,22 @@ DEFAULT_OUTPUT_STYLE_ENVIRONMENT_VARIABLE = EnvironmentVariable( 'COLCON_DEFAULT_OUTPUT_STYLE', 'Select the default output style extension') +# Match CSI (Control Sequence Introducer) and basic G0/G1 sequences. +_ANSI_ESCAPE_RE = re.compile(r'\x1b(?:\[[0-9;]*[a-zA-Z]|[()][a-zA-Z])') + + +def printed_strlen(s): + """ + Get the printed length of a string by skipping over ANSI escape sequences. + + :param s: The string to measure + :returns: The number of printable characters in the string + :rtype: int + """ + if '\x1b' not in s: + return len(s) + return len(_ANSI_ESCAPE_RE.sub('', s)) + class Stylizer(namedtuple('Stylizer', ('start', 'end'))): """A text style modifier.""" diff --git a/test/spell_check.words b/test/spell_check.words index eb55ce8c..b9ad4d09 100644 --- a/test/spell_check.words +++ b/test/spell_check.words @@ -1,5 +1,6 @@ addfinalizer addopts +ansi apache argparse asyncio @@ -63,6 +64,7 @@ hookwrapper https importlib importorskip +introducer isatty isdigit iterdir diff --git a/test/test_output_style.py b/test/test_output_style.py index dc88c1d3..08075f4d 100644 --- a/test/test_output_style.py +++ b/test/test_output_style.py @@ -12,6 +12,7 @@ from colcon_core.output_style import apply_output_style from colcon_core.output_style import DEFAULT_OUTPUT_STYLE_ENVIRONMENT_VARIABLE from colcon_core.output_style import OutputStyleExtensionPoint +from colcon_core.output_style import printed_strlen from colcon_core.output_style import StyleCollection from colcon_core.output_style import Stylizer import pytest @@ -127,3 +128,20 @@ def test_style_dump(mock_stdout): stdout = mock_stdout.getvalue() assert 'loud_errors' in stdout assert 'soft_warnings' in stdout + + +@pytest.mark.parametrize( + 'test_string, expected_length', + [ + ('hello', 5), + ('\x1b[31m' + 'hello' + '\x1b[0m', 5), + ('\x1b[1m' + 'bold' + '\x1b[0m', 4), + ('\x1b[1;31m' + 'red bold' + '\x1b[0m', 8), + ('\x1b[31m' + 'omelet' + '\x1b[0m', 6), + ('\x1b[2K' + 'hello', 5), + ('\x1b[2K' + 'omelet', 6), + ('\x1b(B' + 'hello', 5), + ] +) +def test_printed_strlen(test_string, expected_length): + assert printed_strlen(test_string) == expected_length