From 5b29466c962053c066b97451b5d51209fd0e4019 Mon Sep 17 00:00:00 2001 From: Bob Halley Date: Thu, 13 Aug 2020 17:30:04 -0700 Subject: [PATCH] apply the immutable init wrapper to __setstate__ too, if present --- dns/_immutable_attr.py | 7 +++++++ dns/_immutable_ctx.py | 6 ++++++ 2 files changed, 13 insertions(+) diff --git a/dns/_immutable_attr.py b/dns/_immutable_attr.py index 9d76cdf2..61aa1b17 100644 --- 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 f56864b7..d5766899 100644 --- 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__ -- 2.47.3