_immutable_ctx.py 2.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475
  1. # Copyright (C) Dnspython Contributors, see LICENSE for text of ISC license
  2. # This implementation of the immutable decorator requires python >=
  3. # 3.7, and is significantly more storage efficient when making classes
  4. # with slots immutable. It's also faster.
  5. import contextvars
  6. import inspect
  7. _in__init__ = contextvars.ContextVar('_immutable_in__init__', default=False)
  8. class _Immutable:
  9. """Immutable mixin class"""
  10. # We set slots to the empty list to say "we don't have any attributes".
  11. # We do this so that if we're mixed in with a class with __slots__, we
  12. # don't cause a __dict__ to be added which would waste space.
  13. __slots__ = ()
  14. def __setattr__(self, name, value):
  15. if _in__init__.get() is not self:
  16. raise TypeError("object doesn't support attribute assignment")
  17. else:
  18. super().__setattr__(name, value)
  19. def __delattr__(self, name):
  20. if _in__init__.get() is not self:
  21. raise TypeError("object doesn't support attribute assignment")
  22. else:
  23. super().__delattr__(name)
  24. def _immutable_init(f):
  25. def nf(*args, **kwargs):
  26. previous = _in__init__.set(args[0])
  27. try:
  28. # call the actual __init__
  29. f(*args, **kwargs)
  30. finally:
  31. _in__init__.reset(previous)
  32. nf.__signature__ = inspect.signature(f)
  33. return nf
  34. def immutable(cls):
  35. if _Immutable in cls.__mro__:
  36. # Some ancestor already has the mixin, so just make sure we keep
  37. # following the __init__ protocol.
  38. cls.__init__ = _immutable_init(cls.__init__)
  39. if hasattr(cls, '__setstate__'):
  40. cls.__setstate__ = _immutable_init(cls.__setstate__)
  41. ncls = cls
  42. else:
  43. # Mixin the Immutable class and follow the __init__ protocol.
  44. class ncls(_Immutable, cls):
  45. # We have to do the __slots__ declaration here too!
  46. __slots__ = ()
  47. @_immutable_init
  48. def __init__(self, *args, **kwargs):
  49. super().__init__(*args, **kwargs)
  50. if hasattr(cls, '__setstate__'):
  51. @_immutable_init
  52. def __setstate__(self, *args, **kwargs):
  53. super().__setstate__(*args, **kwargs)
  54. # make ncls have the same name and module as cls
  55. ncls.__name__ = cls.__name__
  56. ncls.__qualname__ = cls.__qualname__
  57. ncls.__module__ = cls.__module__
  58. return ncls