Force unicode before interaction

Handle binary edge case and improve tests

Author: joguSD

awsshell/utils.py

@@ -1,15 +1,18 @@
 """Utility module for misc aws shell functions."""
 from __future__ import print_function
 import os
+import six
 import contextlib
 import tempfile
 import uuid
+import logging
 
 import awscli
 
 from awsshell.compat import HTMLParser
 
 
+LOG = logging.getLogger(__name__)
 AWSCLI_VERSION = awscli.__version__
 
 
@@ -118,3 +121,24 @@ def file_contents(self, filename, binary=False):
 
     def file_exists(self, filename):
         return filename in self._file_mapping
+
+
+def _attempt_decode(string, encoding):
+    try:
+        return string.decode(encoding)
+    except UnicodeError as error:
+        LOG.debug('Failed to decode: %s', error)
+        return string
+
+
+def force_unicode(obj, encoding='utf8'):
+    """Recursively search lists and dicts for strings to encode as unicode."""
+    if isinstance(obj, dict):
+        for key in obj:
+            obj[key] = force_unicode(obj[key])
+    elif isinstance(obj, list):
+        obj = [force_unicode(val) for val in obj]
+    elif isinstance(obj, six.string_types):
+        if not isinstance(obj, six.text_type):
+            obj = _attempt_decode(obj, encoding)
+    return obj

awsshell/wizard.py

@@ -8,6 +8,7 @@
 from botocore.exceptions import BotoCoreError, ClientError
 
 from awsshell.resource import index
+from awsshell.utils import force_unicode
 from awsshell.selectmenu import select_prompt
 from awsshell.interaction import InteractionLoader, InteractionException
 
@@ -275,13 +276,14 @@ def _handle_retrieval(self):
         return data
 
     def _handle_interaction(self, data):
+
         # if no interaction step, just forward data
         if self.interaction is None:
             return data
         else:
             creator = self._interaction_loader.create
             interaction = creator(self.interaction, self.prompt)
-            return interaction.execute(data)
+            return interaction.execute(force_unicode(data))
 
     def _handle_resolution(self, data):
         if self.resolution:

tests/unit/test_utils.py

@@ -2,11 +2,14 @@
 import os
 import tempfile
 import shutil
+import six
+import pytest
 
 from awsshell.utils import FSLayer
 from awsshell.utils import InMemoryFSLayer
 from awsshell.utils import FileReadError
 from awsshell.utils import temporary_file
+from awsshell.utils import force_unicode
 
 
 class TestFSLayer(unittest.TestCase):
@@ -93,3 +96,34 @@ def test_can_open_in_read(self):
             f.seek(0)
             assert f.read() == "foobar"
         self.assertFalse(os.path.isfile(filename))
+
+
+strings = ['', u'', 'some text', u'some text', 'yu\u3086', u'yu\u3086']
+
+
+@pytest.mark.parametrize('text', strings)
+def test_force_unicode(text):
+    unicode_text = force_unicode(text)
+    assert unicode_text == text
+    assert isinstance(unicode_text, six.text_type)
+
+
+def test_force_unicode_binary():
+    # should fail decoding and remain unchanged
+    blob = '\x81'
+    result = force_unicode(blob)
+    assert result is blob
+
+
+def test_force_unicode_recursion():
+    obj = {
+        'a': ['string'],
+        'b': {'str': 'yu\u3086'},
+        'c': '\x81',
+    }
+
+    clean_obj = force_unicode(obj)
+    assert isinstance(clean_obj['a'][0], six.text_type)
+    assert isinstance(clean_obj['b']['str'], six.text_type)
+    assert clean_obj['c'] is obj['c']
+    assert obj == clean_obj