improve netlib.encoding

2 files changed, 38 insertions, 34 deletions

diff --git a/netlib/encoding.py b/netlib/encoding.py
index 9c8acff7..a3c83c46 100644
--- a/netlib/encoding.py
+++ b/netlib/encoding.py
@@ -5,7 +5,9 @@ from __future__ import absolute_import
 
 import codecs
 import collections
+import six
 from io import BytesIO
+
 import gzip
 import zlib
 import brotli
@@ -32,6 +34,9 @@ def decode(encoded, encoding, errors='strict'):
     Raises:
         ValueError, if decoding fails.
     """
+    if len(encoded) == 0:
+        return encoded
+
     global _cache
     cached = (
         isinstance(encoded, bytes) and
@@ -49,11 +54,14 @@ def decode(encoded, encoding, errors='strict'):
         if encoding in ("gzip", "deflate", "br"):
             _cache = CachedDecode(encoded, encoding, errors, decoded)
         return decoded
+    except TypeError:
+        raise
     except Exception as e:
-        raise ValueError("{} when decoding {} with {}".format(
+        raise ValueError("{} when decoding {} with {}: {}".format(
             type(e).__name__,
             repr(encoded)[:10],
             repr(encoding),
+            repr(e),
         ))
 
 
@@ -68,6 +76,9 @@ def encode(decoded, encoding, errors='strict'):
     Raises:
         ValueError, if encoding fails.
     """
+    if len(decoded) == 0:
+        return decoded
+
     global _cache
     cached = (
         isinstance(decoded, bytes) and
@@ -79,17 +90,23 @@ def encode(decoded, encoding, errors='strict'):
         return _cache.encoded
     try:
         try:
-            encoded = custom_encode[encoding](decoded)
+            value = decoded
+            if not six.PY2 and isinstance(value, six.string_types):
+                value = decoded.encode()
+            encoded = custom_encode[encoding](value)
         except KeyError:
             encoded = codecs.encode(decoded, encoding, errors)
         if encoding in ("gzip", "deflate", "br"):
             _cache = CachedDecode(encoded, encoding, errors, decoded)
         return encoded
+    except TypeError:
+        raise
     except Exception as e:
-        raise ValueError("{} when encoding {} with {}".format(
+        raise ValueError("{} when encoding {} with {}: {}".format(
             type(e).__name__,
             repr(decoded)[:10],
             repr(encoding),
+            repr(e),
         ))
 
 
diff --git a/test/netlib/test_encoding.py b/test/netlib/test_encoding.py
index 08e69ec5..b191d50d 100644
--- a/test/netlib/test_encoding.py
+++ b/test/netlib/test_encoding.py
@@ -1,4 +1,6 @@
 import mock
+import pytest
+
 from netlib import encoding, tutils
 
 
@@ -9,47 +11,32 @@ def test_identity():
         encoding.encode(b"string", "nonexistent encoding")
 
 
-def test_gzip():
-    assert b"string" == encoding.decode(
-        encoding.encode(
-            b"string",
-            "gzip"
-        ),
-        "gzip"
-    )
-    with tutils.raises(ValueError):
-        encoding.decode(b"bogus", "gzip")
+@pytest.mark.parametrize("encoder", [
+    'gzip',
+    'br',
+    'deflate',
+])
+def test_encoders(encoder):
+    assert "" == encoding.decode("", encoder)
+    assert b"" == encoding.decode(b"", encoder)
 
-
-def test_brotli():
-    assert b"string" == encoding.decode(
+    assert "string" == encoding.decode(
         encoding.encode(
-            b"string",
-            "br"
+            "string",
+            encoder
         ),
-        "br"
+        encoder
     )
-    with tutils.raises(ValueError):
-        encoding.decode(b"bogus", "br")
-
-
-def test_deflate():
     assert b"string" == encoding.decode(
         encoding.encode(
             b"string",
-            "deflate"
+            encoder
         ),
-        "deflate"
-    )
-    assert b"string" == encoding.decode(
-        encoding.encode(
-            b"string",
-            "deflate"
-        )[2:-4],
-        "deflate"
+        encoder
     )
+
     with tutils.raises(ValueError):
-        encoding.decode(b"bogus", "deflate")
+        encoding.decode(b"bogus", encoder)
 
 
 def test_cache():