fix py3 compat

5 files changed, 10 insertions, 9 deletions

diff --git a/netlib/certutils.py b/netlib/certutils.py
index ecdc0624..616a778e 100644
--- a/netlib/certutils.py
+++ b/netlib/certutils.py
@@ -12,8 +12,9 @@ from pyasn1.codec.der.decoder import decode
 from pyasn1.error import PyAsn1Error
 import OpenSSL
 
+from .utils import Serializable
+
 # Default expiry must not be too long: https://github.com/mitmproxy/mitmproxy/issues/815
-from netlib.utils import Serializable
 
 DEFAULT_EXP = 94608000  # = 24 * 60 * 60 * 365 * 3
 # Generated with "openssl dhparam". It's too slow to generate this on startup.
diff --git a/netlib/http/message.py b/netlib/http/message.py
index 3d65f93e..e3d8ce37 100644
--- a/netlib/http/message.py
+++ b/netlib/http/message.py
@@ -4,7 +4,6 @@ import warnings
 
 import six
 
-from netlib.utils import Serializable
 from .headers import Headers
 from .. import encoding, utils
 
@@ -19,7 +18,7 @@ else:
     _always_bytes = lambda x: utils.always_bytes(x, "utf-8", "surrogateescape")
 
 
-class MessageData(Serializable):
+class MessageData(utils.Serializable):
     def __eq__(self, other):
         if isinstance(other, MessageData):
             return self.__dict__ == other.__dict__
@@ -45,7 +44,7 @@ class MessageData(Serializable):
         return cls(**state)
 
 
-class Message(Serializable):
+class Message(utils.Serializable):
     def __init__(self, data):
         self.data = data
 
diff --git a/netlib/tcp.py b/netlib/tcp.py
index 2e91a70c..c8548aea 100644
--- a/netlib/tcp.py
+++ b/netlib/tcp.py
@@ -16,8 +16,7 @@ import six
 import OpenSSL
 from OpenSSL import SSL
 
-from netlib.utils import Serializable
-from . import certutils, version_check
+from . import certutils, version_check, utils
 
 # This is a rather hackish way to make sure that
 # the latest version of pyOpenSSL is actually installed.
@@ -299,7 +298,7 @@ class Reader(_FileLike):
             raise NotImplementedError("Can only peek into (pyOpenSSL) sockets")
 
 
-class Address(Serializable):
+class Address(utils.Serializable):
 
     """
         This class wraps an IPv4/IPv6 tuple to provide named attributes and
diff --git a/test/http/test_request.py b/test/http/test_request.py
index 1deee387..900b2cd1 100644
--- a/test/http/test_request.py
+++ b/test/http/test_request.py
@@ -12,7 +12,7 @@ from .test_message import _test_decoded_attr, _test_passthrough_attr
 
 class TestRequestData(object):
     def test_init(self):
-        with raises(ValueError):
+        with raises(ValueError if six.PY2 else TypeError):
             treq(headers="foobar")
 
         assert isinstance(treq(headers=None).headers, Headers)
diff --git a/test/http/test_response.py b/test/http/test_response.py
index c7d90b16..14588000 100644
--- a/test/http/test_response.py
+++ b/test/http/test_response.py
@@ -1,5 +1,7 @@
 from __future__ import absolute_import, print_function, division
 
+import six
+
 from netlib.http import Headers
 from netlib.odict import ODict, ODictCaseless
 from netlib.tutils import raises, tresp
@@ -8,7 +10,7 @@ from .test_message import _test_passthrough_attr, _test_decoded_attr
 
 class TestResponseData(object):
     def test_init(self):
-        with raises(ValueError):
+        with raises(ValueError if six.PY2 else TypeError):
             tresp(headers="foobar")
 
         assert isinstance(tresp(headers=None).headers, Headers)