_threadsafety.py 1.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758
  1. import threading
  2. import scipy._lib.decorator
  3. __all__ = ['ReentrancyError', 'ReentrancyLock', 'non_reentrant']
  4. class ReentrancyError(RuntimeError):
  5. pass
  6. class ReentrancyLock:
  7. """
  8. Threading lock that raises an exception for reentrant calls.
  9. Calls from different threads are serialized, and nested calls from the
  10. same thread result to an error.
  11. The object can be used as a context manager or to decorate functions
  12. via the decorate() method.
  13. """
  14. def __init__(self, err_msg):
  15. self._rlock = threading.RLock()
  16. self._entered = False
  17. self._err_msg = err_msg
  18. def __enter__(self):
  19. self._rlock.acquire()
  20. if self._entered:
  21. self._rlock.release()
  22. raise ReentrancyError(self._err_msg)
  23. self._entered = True
  24. def __exit__(self, type, value, traceback):
  25. self._entered = False
  26. self._rlock.release()
  27. def decorate(self, func):
  28. def caller(func, *a, **kw):
  29. with self:
  30. return func(*a, **kw)
  31. return scipy._lib.decorator.decorate(func, caller)
  32. def non_reentrant(err_msg=None):
  33. """
  34. Decorate a function with a threading lock and prevent reentrant calls.
  35. """
  36. def decorator(func):
  37. msg = err_msg
  38. if msg is None:
  39. msg = "%s is not re-entrant" % func.__name__
  40. lock = ReentrancyLock(msg)
  41. return lock.decorate(func)
  42. return decorator