Wait until later to convert to byte strings for output

Also some pre-emptive python3 compat and a code simplification

1 files changed, 50 insertions, 26 deletions

diff --git a/lib/ansible/utils/display.py b/lib/ansible/utils/display.py
index 63f5fc4c53..bc21820033 100644
--- a/lib/ansible/utils/display.py
+++ b/lib/ansible/utils/display.py
@@ -35,8 +35,14 @@ from multiprocessing import Lock
 from ansible import constants as C
 from ansible.errors import AnsibleError
 from ansible.utils.color import stringc
-from ansible.utils.unicode import to_bytes
+from ansible.utils.unicode import to_bytes, to_unicode
 
+try:
+    # Python 2
+    input = raw_input
+except NameError:
+    # Python 3
+    pass
 
 
 # These are module level as we currently fork and serialize the whole process and locks in the objects don't play well with that
@@ -96,30 +102,50 @@ class Display:
                 self.noncow = random.choice(cows)
 
     def display(self, msg, color=None, stderr=False, screen_only=False, log_only=False):
+        """ Display a message to the user
+
+        Note: msg *must* be a unicode string to prevent UnicodeError tracebacks.
+        """ 
 
         # FIXME: this needs to be implemented
         #msg = utils.sanitize_output(msg)
-        msg2 = self._safe_output(msg, stderr=stderr)
         if color:
-            msg2 = stringc(msg, color)
+            msg = stringc(msg, color)
 
         if not log_only:
-            b_msg2 = to_bytes(msg2)
+            if not msg.endswith(u'\n'):
+                msg2 = msg + u'\n'
+            else:
+                msg2 = msg
+
+            msg2 = to_bytes(msg2, encoding=self._output_encoding(stderr=stderr))
+            if sys.version_info >= (3,):
+                # Convert back to text string on python3
+                # We first convert to a byte string so that we get rid of
+                # characters that are invalid in the user's locale
+                msg2 = to_unicode(msg2, self._output_encoding(stderr=stderr))
+
             if not stderr:
-                print(b_msg2)
+                sys.stdout.write(msg2)
                 sys.stdout.flush()
             else:
-                print(b_msg2, file=sys.stderr)
+                sys.stderr.write(msg2)
                 sys.stderr.flush()
 
         if logger and not screen_only:
-            while msg.startswith("\n"):
-                msg = msg.replace("\n","")
-            b_msg = to_bytes(msg)
+            msg2 = msg.lstrip(u'\n')
+
+            msg2 = to_bytes(msg2)
+            if sys.version_info >= (3,):
+                # Convert back to text string on python3
+                # We first convert to a byte string so that we get rid of
+                # characters that are invalid in the user's locale
+                msg2 = to_unicode(msg2, self._output_encoding(stderr=stderr))
+
             if color == 'red':
-                logger.error(b_msg)
+                logger.error(msg2)
             else:
-                logger.info(b_msg)
+                logger.info(msg2)
 
     def vv(self, msg, host=None):
         return self.verbose(msg, host=host, caplevel=1)
@@ -230,21 +256,20 @@ class Display:
             self.display(new_msg, color='red', stderr=True)
             self._errors[new_msg] = 1
 
+    @staticmethod
     def prompt(self, msg):
-
-        return raw_input(self._safe_output(msg))
-
-    def _safe_output(self, msg, stderr=False):
-
-        encoding='utf-8'
-        if not stderr and sys.stdout.encoding:
-            encoding = sys.stdout.encoding
-        elif stderr and sys.stderr.encoding:
-            encoding = sys.stderr.encoding
-
-        msg = to_bytes(msg, encoding)
-
-        return msg
+        prompt_string = to_bytes(msg, encoding=self._output_encoding())
+        if sys.version_info >= (3,):
+            # Convert back into text on python3.  We do this double conversion
+            # to get rid of characters that are illegal in the user's locale
+            prompt_string = to_unicode(prompt_string)
+        return input(prompt_string)
+
+    @staticmethod
+    def _output_encoding(stderr=False):
+        if stderr:
+            return sys.stderr.encoding or 'utf-8'
+        return sys.stout.encoding or 'utf-8'
 
     def _set_column_width(self):
         if os.isatty(0):
@@ -252,4 +277,3 @@ class Display:
         else:
             tty_size = 0
         self.columns = max(79, tty_size)
-