diff options
-rw-r--r-- | dns/_immutable_attr.py | 2 | ||||
-rw-r--r-- | dns/_immutable_ctx.py | 2 |
2 files changed, 4 insertions, 0 deletions
diff --git a/dns/_immutable_attr.py b/dns/_immutable_attr.py index 61aa1b1..bf8aebc 100644 --- a/dns/_immutable_attr.py +++ b/dns/_immutable_attr.py @@ -56,6 +56,8 @@ def immutable(cls): # Some ancestor already has the mixin, so just make sure we keep # following the __init__ protocol. cls.__init__ = _immutable_init(cls.__init__) + if hasattr(cls, '__setstate__'): + cls.__setstate__ = _immutable_init(cls.__setstate__) ncls = cls else: # Mixin the Immutable class and follow the __init__ protocol. diff --git a/dns/_immutable_ctx.py b/dns/_immutable_ctx.py index d576689..400ab6c 100644 --- a/dns/_immutable_ctx.py +++ b/dns/_immutable_ctx.py @@ -47,6 +47,8 @@ def immutable(cls): # Some ancestor already has the mixin, so just make sure we keep # following the __init__ protocol. cls.__init__ = _immutable_init(cls.__init__) + if hasattr(cls, '__setstate__'): + cls.__setstate__ = _immutable_init(cls.__setstate__) ncls = cls else: # Mixin the Immutable class and follow the __init__ protocol. |