serialization.py 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164
  1. # -*- coding: utf-8 -*-
  2. """
  3. celery.utils.serialization
  4. ~~~~~~~~~~~~~~~~~~~~~~~~~~
  5. Utilities for safely pickling exceptions.
  6. """
  7. from __future__ import absolute_import
  8. from inspect import getmro
  9. from itertools import takewhile
  10. try:
  11. import cPickle as pickle
  12. except ImportError:
  13. import pickle # noqa
  14. from .encoding import safe_repr
  15. __all__ = ['UnpickleableExceptionWrapper', 'subclass_exception',
  16. 'find_pickleable_exception', 'create_exception_cls',
  17. 'get_pickleable_exception', 'get_pickleable_etype',
  18. 'get_pickled_exception']
  19. #: List of base classes we probably don't want to reduce to.
  20. try:
  21. unwanted_base_classes = (StandardError, Exception, BaseException, object)
  22. except NameError: # pragma: no cover
  23. unwanted_base_classes = (Exception, BaseException, object) # py3k
  24. def subclass_exception(name, parent, module): # noqa
  25. return type(name, (parent, ), {'__module__': module})
  26. def find_pickleable_exception(exc, loads=pickle.loads,
  27. dumps=pickle.dumps):
  28. """With an exception instance, iterate over its super classes (by mro)
  29. and find the first super exception that is pickleable. It does
  30. not go below :exc:`Exception` (i.e. it skips :exc:`Exception`,
  31. :class:`BaseException` and :class:`object`). If that happens
  32. you should use :exc:`UnpickleableException` instead.
  33. :param exc: An exception instance.
  34. Will return the nearest pickleable parent exception class
  35. (except :exc:`Exception` and parents), or if the exception is
  36. pickleable it will return :const:`None`.
  37. :rtype :exc:`Exception`:
  38. """
  39. exc_args = getattr(exc, 'args', [])
  40. for supercls in itermro(exc.__class__, unwanted_base_classes):
  41. try:
  42. superexc = supercls(*exc_args)
  43. loads(dumps(superexc))
  44. except:
  45. pass
  46. else:
  47. return superexc
  48. find_nearest_pickleable_exception = find_pickleable_exception # XXX compat
  49. def itermro(cls, stop):
  50. return takewhile(lambda sup: sup not in stop, getmro(cls))
  51. def create_exception_cls(name, module, parent=None):
  52. """Dynamically create an exception class."""
  53. if not parent:
  54. parent = Exception
  55. return subclass_exception(name, parent, module)
  56. class UnpickleableExceptionWrapper(Exception):
  57. """Wraps unpickleable exceptions.
  58. :param exc_module: see :attr:`exc_module`.
  59. :param exc_cls_name: see :attr:`exc_cls_name`.
  60. :param exc_args: see :attr:`exc_args`
  61. **Example**
  62. .. code-block:: python
  63. >>> try:
  64. ... something_raising_unpickleable_exc()
  65. >>> except Exception as e:
  66. ... exc = UnpickleableException(e.__class__.__module__,
  67. ... e.__class__.__name__,
  68. ... e.args)
  69. ... pickle.dumps(exc) # Works fine.
  70. """
  71. #: The module of the original exception.
  72. exc_module = None
  73. #: The name of the original exception class.
  74. exc_cls_name = None
  75. #: The arguments for the original exception.
  76. exc_args = None
  77. def __init__(self, exc_module, exc_cls_name, exc_args, text=None):
  78. safe_exc_args = []
  79. for arg in exc_args:
  80. try:
  81. pickle.dumps(arg)
  82. safe_exc_args.append(arg)
  83. except Exception:
  84. safe_exc_args.append(safe_repr(arg))
  85. self.exc_module = exc_module
  86. self.exc_cls_name = exc_cls_name
  87. self.exc_args = safe_exc_args
  88. self.text = text
  89. Exception.__init__(self, exc_module, exc_cls_name, safe_exc_args, text)
  90. def restore(self):
  91. return create_exception_cls(self.exc_cls_name,
  92. self.exc_module)(*self.exc_args)
  93. def __str__(self):
  94. return self.text
  95. @classmethod
  96. def from_exception(cls, exc):
  97. return cls(exc.__class__.__module__,
  98. exc.__class__.__name__,
  99. getattr(exc, 'args', []),
  100. safe_repr(exc))
  101. def get_pickleable_exception(exc):
  102. """Make sure exception is pickleable."""
  103. try:
  104. pickle.loads(pickle.dumps(exc))
  105. except Exception:
  106. pass
  107. else:
  108. return exc
  109. nearest = find_pickleable_exception(exc)
  110. if nearest:
  111. return nearest
  112. return UnpickleableExceptionWrapper.from_exception(exc)
  113. def get_pickleable_etype(cls, loads=pickle.loads, dumps=pickle.dumps):
  114. try:
  115. loads(dumps(cls))
  116. except:
  117. return Exception
  118. else:
  119. return cls
  120. def get_pickled_exception(exc):
  121. """Get original exception from exception pickled using
  122. :meth:`get_pickleable_exception`."""
  123. if isinstance(exc, UnpickleableExceptionWrapper):
  124. return exc.restore()
  125. return exc