Você não pode selecionar mais de 25 tópicos Os tópicos devem começar com uma letra ou um número, podem incluir traços ('-') e podem ter até 35 caracteres.
 
 
 
 

40 linhas
1.4 KiB

  1. from contextlib import asynccontextmanager as asynccontextmanager
  2. from typing import AsyncGenerator, ContextManager, TypeVar
  3. import anyio.to_thread
  4. from anyio import CapacityLimiter
  5. from starlette.concurrency import iterate_in_threadpool as iterate_in_threadpool # noqa
  6. from starlette.concurrency import run_in_threadpool as run_in_threadpool # noqa
  7. from starlette.concurrency import ( # noqa
  8. run_until_first_complete as run_until_first_complete,
  9. )
  10. _T = TypeVar("_T")
  11. @asynccontextmanager
  12. async def contextmanager_in_threadpool(
  13. cm: ContextManager[_T],
  14. ) -> AsyncGenerator[_T, None]:
  15. # blocking __exit__ from running waiting on a free thread
  16. # can create race conditions/deadlocks if the context manager itself
  17. # has its own internal pool (e.g. a database connection pool)
  18. # to avoid this we let __exit__ run without a capacity limit
  19. # since we're creating a new limiter for each call, any non-zero limit
  20. # works (1 is arbitrary)
  21. exit_limiter = CapacityLimiter(1)
  22. try:
  23. yield await run_in_threadpool(cm.__enter__)
  24. except Exception as e:
  25. ok = bool(
  26. await anyio.to_thread.run_sync(
  27. cm.__exit__, type(e), e, e.__traceback__, limiter=exit_limiter
  28. )
  29. )
  30. if not ok:
  31. raise e
  32. else:
  33. await anyio.to_thread.run_sync(
  34. cm.__exit__, None, None, None, limiter=exit_limiter
  35. )