Move the encoding result check to correct place

Author: msimacek

graalpython/lib-graalpython/_codecs.py

@@ -121,11 +121,7 @@ def encode(obj, encoding='utf-8', errors='strict'):
         result = encoder(obj, errors)
         if not (isinstance(result, tuple) and len(result) == 2):
             raise TypeError('encoder must return a tuple (object, integer)')
-        encoded = result[0]
-        if not isinstance(encoded, bytes):
-            raise TypeError("'%s' encoder returned '%s' instead of 'bytes'; use codecs.encode() to encode to arbitrary types"
-                            % (encoding, type(encoded).__name__))
-        return encoded
+        return result[0]
 
 
 @__graalpython__.builtin
@@ -135,11 +131,7 @@ def decode(obj, encoding='utf-8', errors='strict'):
         result = decoder(obj, errors)
         if not (isinstance(result, tuple) and len(result) == 2):
             raise TypeError('decoder must return a tuple (object, integer)')
-        decoded = result[0]
-        if not isinstance(decoded, str):
-            raise TypeError("'%s' encoder returned '%s' instead of 'str'; use codecs.encode() to encode to arbitrary types"
-                            % (encoding, type(decoded).__name__))
-        return decoded
+        return result[0]
 
 
 @__graalpython__.builtin

graalpython/lib-graalpython/bytes.py

@@ -52,7 +52,11 @@ def decode(self, encoding="utf-8", errors="strict"):
       as well as any other name registered with codecs.register_error that
       can handle UnicodeDecodeErrors.
     """
-    return _codecs.decode(self, encoding=encoding, errors=errors)
+    result = _codecs.decode(self, encoding=encoding, errors=errors)
+    if not isinstance(result, str):
+        raise TypeError("'%s' encoder returned '%s' instead of 'str'; use codecs.decode() to decode to arbitrary types"
+                        % (encoding, type(result).__name__))
+    return result
 
 
 bytes.decode = __graalpython__.builtin_method(decode)

graalpython/lib-graalpython/str.py

@@ -375,7 +375,13 @@ def encode(self, encoding="utf-8", errors="strict"):
       as well as any other name registered with codecs.register_error that
       can handle UnicodeDecodeErrors.
     """
-    return _codecs.encode(self, encoding=encoding, errors=errors)
+    result = _codecs.encode(self, encoding=encoding, errors=errors)
+    if not isinstance(result, bytes):
+        if isinstance(result, bytearray):
+            return bytes(result)
+        raise TypeError("'%s' encoder returned '%s' instead of 'bytes'; use codecs.encode() to encode to arbitrary types"
+                        % (encoding, type(result).__name__))
+    return result
 
 
 str.encode = encode