apply the immutable init wrapper to __setstate__ too, if present

diff --git a/dns/_immutable_attr.py b/dns/_immutable_attr.py
index 9d76cdf26fea686a90f9d3dfbacf300793484706..61aa1b173a09cb0b94a9437895ec120861b77188 100644 (file)
--- a/dns/_immutable_attr.py
+++ b/dns/_immutable_attr.py
@@ -60,9 +60,16 @@ def immutable(cls):
     else:
         # Mixin the Immutable class and follow the __init__ protocol.
         class ncls(_Immutable, cls):
+
             @_immutable_init
             def __init__(self, *args, **kwargs):
                 super().__init__(*args, **kwargs)
+
+            if hasattr(cls, '__setstate__'):
+                @_immutable_init
+                def __setstate__(self, *args, **kwargs):
+                    super().__setstate__(*args, **kwargs)
+
         # make ncls have the same name and module as cls
         ncls.__name__ = cls.__name__
         ncls.__qualname__ = cls.__qualname__

diff --git a/dns/_immutable_ctx.py b/dns/_immutable_ctx.py
index f56864b7e3cb2c11a81c91fe63a488a492f874a0..d57668999655d6246ed4b44aa917016febf30ee1 100644 (file)
--- a/dns/_immutable_ctx.py
+++ b/dns/_immutable_ctx.py
@@ -57,6 +57,12 @@ def immutable(cls):
             @_immutable_init
             def __init__(self, *args, **kwargs):
                 super().__init__(*args, **kwargs)
+
+            if hasattr(cls, '__setstate__'):
+                @_immutable_init
+                def __setstate__(self, *args, **kwargs):
+                    super().__setstate__(*args, **kwargs)
+
         # make ncls have the same name and module as cls
         ncls.__name__ = cls.__name__
         ncls.__qualname__ = cls.__qualname__