Use libc wcwidth to calculate print width in display (#66214)

* Use libc wcwidth to calculate print width in display. Fixes #63105

* Remove errantly added blank lines

* Fixes

* Move setlocale, adjust tests to work around py2 oddity with characters following null

* Don't change cli stub

* emojis

* Remove to_text call

* Special accounting for deletions

* Add initialization function, expand tests, ensure fallback to len

* get_text_width requires text, ensure banner deals with it

* Handle setlocale errors

* Move variable decrement

* Remove unused import

changelogs/fragments/63105-wcswidth.yml

@@ -0,0 +1,3 @@
+bugfixes:
+- Display - Use wcswidth to calculate printable width of a text string
+  (https://github.com/ansible/ansible/issues/63105)

lib/ansible/cli/scripts/ansible_cli_stub.py

@@ -60,11 +60,13 @@ if __name__ == '__main__':
 
     try:  # bad ANSIBLE_CONFIG or config options can force ugly stacktrace
         import ansible.constants as C
-        from ansible.utils.display import Display
+        from ansible.utils.display import Display, initialize_locale
     except AnsibleOptionsError as e:
         display.error(to_text(e), wrap_text=False)
         sys.exit(5)
 
+    initialize_locale()
+
     cli = None
     me = os.path.basename(sys.argv[0])
 

lib/ansible/utils/display.py

@@ -18,6 +18,7 @@
 from __future__ import (absolute_import, division, print_function)
 __metaclass__ = type
 
+import ctypes.util
 import errno
 import fcntl
 import getpass
@@ -36,7 +37,7 @@ from termios import TIOCGWINSZ
 from ansible import constants as C
 from ansible.errors import AnsibleError, AnsibleAssertionError
 from ansible.module_utils._text import to_bytes, to_text, to_native
-from ansible.module_utils.six import with_metaclass, string_types
+from ansible.module_utils.six import with_metaclass, text_type
 from ansible.utils.color import stringc
 from ansible.utils.singleton import Singleton
 from ansible.utils.unsafe_proxy import wrap_var
@@ -48,6 +49,100 @@ except NameError:
     # Python 3, we already have raw_input
     pass
 
+_LIBC = ctypes.cdll.LoadLibrary(ctypes.util.find_library('c'))
+# Set argtypes, to avoid segfault if the wrong type is provided,
+# restype is assumed to be c_int
+_LIBC.wcwidth.argtypes = (ctypes.c_wchar,)
+_LIBC.wcswidth.argtypes = (ctypes.c_wchar_p, ctypes.c_int)
+# Max for c_int
+_MAX_INT = 2 ** (ctypes.sizeof(ctypes.c_int) * 8 - 1) - 1
+
+_LOCALE_INITIALIZED = False
+_LOCALE_INITIALIZATION_ERR = None
+
+
+def initialize_locale():
+    """Set the locale to the users default setting
+    and set ``_LOCALE_INITIALIZED`` to indicate whether
+    ``get_text_width`` may run into trouble
+    """
+    global _LOCALE_INITIALIZED, _LOCALE_INITIALIZATION_ERR
+    if _LOCALE_INITIALIZED is False:
+        try:
+            locale.setlocale(locale.LC_ALL, '')
+        except locale.Error as e:
+            _LOCALE_INITIALIZATION_ERR = e
+        else:
+            _LOCALE_INITIALIZED = True
+
+
+def get_text_width(text):
+    """Function that utilizes ``wcswidth`` or ``wcwidth`` to determine the
+    number of columns used to display a text string.
+
+    We try first with ``wcswidth``, and fallback to iterating each
+    character and using wcwidth individually, falling back to a value of 0
+    for non-printable wide characters
+
+    On Py2, this depends on ``locale.setlocale(locale.LC_ALL, '')``,
+    that in the case of Ansible is done in ``bin/ansible``
+    """
+    if not isinstance(text, text_type):
+        raise TypeError('get_text_width requires text, not %s' % type(text))
+
+    if _LOCALE_INITIALIZATION_ERR:
+        Display().warning(
+            'An error occurred while calling ansible.utils.display.initialize_locale '
+            '(%s). This may result in incorrectly calculated text widths that can '
+            'cause Display to print incorrect line lengths' % _LOCALE_INITIALIZATION_ERR
+        )
+    elif not _LOCALE_INITIALIZED:
+        Display().warning(
+            'ansible.utils.display.initialize_locale has not been called, '
+            'this may result in incorrectly calculated text widths that can '
+            'cause Display to print incorrect line lengths'
+        )
+
+    try:
+        width = _LIBC.wcswidth(text, _MAX_INT)
+    except ctypes.ArgumentError:
+        width = -1
+    if width != -1:
+        return width
+
+    width = 0
+    counter = 0
+    for c in text:
+        counter += 1
+        if c in (u'\x08', u'\x7f', u'\x94', u'\x1b'):
+            # A few characters result in a subtraction of length:
+            # BS, DEL, CCH, ESC
+            # ESC is slightly different in that it's part of an escape sequence, and
+            # while ESC is non printable, it's part of an escape sequence, which results
+            # in a single non printable length
+            width -= 1
+            counter -= 1
+            continue
+
+        try:
+            w = _LIBC.wcwidth(c)
+        except ctypes.ArgumentError:
+            w = -1
+        if w == -1:
+            # -1 signifies a non-printable character
+            # use 0 here as a best effort
+            w = 0
+        width += w
+
+    if width == 0 and counter and not _LOCALE_INITIALIZED:
+        raise EnvironmentError(
+            'ansible.utils.display.initialize_locale has not been called, '
+            'and get_text_width could not calculate text width of %r' % text
+        )
+
+    # It doesn't make sense to have a negative printable width
+    return width if width >= 0 else 0
+
 
 class FilterBlackList(logging.Filter):
     def __init__(self, blacklist):
@@ -321,6 +416,8 @@ class Display(with_metaclass(Singleton, object)):
         '''
         Prints a header-looking line with cowsay or stars with length depending on terminal width (3 minimum)
         '''
+        msg = to_text(msg)
+
         if self.b_cowsay and cows:
             try:
                 self.banner_cowsay(msg)
@@ -329,7 +426,10 @@ class Display(with_metaclass(Singleton, object)):
                 self.warning("somebody cleverly deleted cowsay or something during the PB run.  heh.")
 
         msg = msg.strip()
-        star_len = self.columns - len(msg)
+        try:
+            star_len = self.columns - get_text_width(msg)
+        except EnvironmentError:
+            star_len = self.columns - len(msg)
         if star_len <= 3:
             star_len = 3
         stars = u"*" * star_len

test/units/utils/test_display.py

@@ -0,0 +1,65 @@
+# -*- coding: utf-8 -*-
+# (c) 2020 Matt Martz <matt@sivel.net>
+# GNU General Public License v3.0+ (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt)
+
+from __future__ import absolute_import, division, print_function
+__metaclass__ = type
+
+from units.compat.mock import MagicMock
+
+import pytest
+
+from ansible.module_utils.six import PY3
+from ansible.utils.display import Display, get_text_width, initialize_locale
+
+
+def test_get_text_width():
+    initialize_locale()
+    assert get_text_width(u'コンニチハ') == 10
+    assert get_text_width(u'abコcd') == 6
+    assert get_text_width(u'café') == 4
+    assert get_text_width(u'four') == 4
+    assert get_text_width(u'\u001B') == 0
+    assert get_text_width(u'ab\u0000') == 2
+    assert get_text_width(u'abコ\u0000') == 4
+    assert get_text_width(u'🚀🐮') == 4
+    assert get_text_width(u'\x08') == 0
+    assert get_text_width(u'\x08\x08') == 0
+    assert get_text_width(u'ab\x08cd') == 3
+    assert get_text_width(u'ab\x1bcd') == 3
+    assert get_text_width(u'ab\x7fcd') == 3
+    assert get_text_width(u'ab\x94cd') == 3
+
+    pytest.raises(TypeError, get_text_width, 1)
+    pytest.raises(TypeError, get_text_width, b'four')
+
+
+@pytest.mark.skipif(PY3, reason='Fallback only happens reliably on py2')
+def test_get_text_width_no_locale():
+    pytest.raises(EnvironmentError, get_text_width, u'🚀🐮')
+
+
+def test_Display_banner_get_text_width(monkeypatch):
+    initialize_locale()
+    display = Display()
+    display_mock = MagicMock()
+    monkeypatch.setattr(display, 'display', display_mock)
+
+    display.banner(u'🚀🐮', color=False, cows=False)
+    args, kwargs = display_mock.call_args
+    msg = args[0]
+    stars = u' %s' % (75 * u'*')
+    assert msg.endswith(stars)
+
+
+@pytest.mark.skipif(PY3, reason='Fallback only happens reliably on py2')
+def test_Display_banner_get_text_width_fallback(monkeypatch):
+    display = Display()
+    display_mock = MagicMock()
+    monkeypatch.setattr(display, 'display', display_mock)
+
+    display.banner(u'🚀🐮', color=False, cows=False)
+    args, kwargs = display_mock.call_args
+    msg = args[0]
+    stars = u' %s' % (77 * u'*')
+    assert msg.endswith(stars)