asyncio.py 1.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364
  1. import os
  2. from asyncio import get_running_loop
  3. from functools import wraps
  4. from django.core.exceptions import SynchronousOnlyOperation
  5. def async_unsafe(message):
  6. """
  7. Decorator to mark functions as async-unsafe. Someone trying to access
  8. the function while in an async context will get an error message.
  9. """
  10. def decorator(func):
  11. @wraps(func)
  12. def inner(*args, **kwargs):
  13. # Detect a running event loop in this thread.
  14. try:
  15. get_running_loop()
  16. except RuntimeError:
  17. pass
  18. else:
  19. if not os.environ.get("DJANGO_ALLOW_ASYNC_UNSAFE"):
  20. raise SynchronousOnlyOperation(message)
  21. # Pass onward.
  22. return func(*args, **kwargs)
  23. return inner
  24. # If the message is actually a function, then be a no-arguments decorator.
  25. if callable(message):
  26. func = message
  27. message = (
  28. "You cannot call this from an async context - use a thread or "
  29. "sync_to_async."
  30. )
  31. return decorator(func)
  32. else:
  33. return decorator
  34. try:
  35. from contextlib import aclosing
  36. except ImportError:
  37. # TODO: Remove when dropping support for PY39.
  38. from contextlib import AbstractAsyncContextManager
  39. # Backport of contextlib.aclosing() from Python 3.10. Copyright (C) Python
  40. # Software Foundation (see LICENSE.python).
  41. class aclosing(AbstractAsyncContextManager):
  42. """
  43. Async context manager for safely finalizing an asynchronously
  44. cleaned-up resource such as an async generator, calling its
  45. ``aclose()`` method.
  46. """
  47. def __init__(self, thing):
  48. self.thing = thing
  49. async def __aenter__(self):
  50. return self.thing
  51. async def __aexit__(self, *exc_info):
  52. await self.thing.aclose()