concurrency.py 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110
  1. # util/concurrency.py
  2. # Copyright (C) 2005-2026 the SQLAlchemy authors and contributors
  3. # <see AUTHORS file>
  4. #
  5. # This module is part of SQLAlchemy and is released under
  6. # the MIT License: https://www.opensource.org/licenses/mit-license.php
  7. # mypy: allow-untyped-defs, allow-untyped-calls
  8. """asyncio-related concurrency functions."""
  9. from __future__ import annotations
  10. import asyncio # noqa
  11. import typing
  12. from typing import Any
  13. from typing import Callable
  14. from typing import Coroutine
  15. from typing import TypeVar
  16. have_greenlet = False
  17. greenlet_error = None
  18. try:
  19. import greenlet # type: ignore[import-untyped,unused-ignore] # noqa: F401,E501
  20. except ImportError as e:
  21. greenlet_error = str(e)
  22. pass
  23. else:
  24. have_greenlet = True
  25. from ._concurrency_py3k import await_only as await_only
  26. from ._concurrency_py3k import await_fallback as await_fallback
  27. from ._concurrency_py3k import in_greenlet as in_greenlet
  28. from ._concurrency_py3k import greenlet_spawn as greenlet_spawn
  29. from ._concurrency_py3k import is_exit_exception as is_exit_exception
  30. from ._concurrency_py3k import AsyncAdaptedLock as AsyncAdaptedLock
  31. from ._concurrency_py3k import _Runner
  32. _T = TypeVar("_T")
  33. class _AsyncUtil:
  34. """Asyncio util for test suite/ util only"""
  35. def __init__(self) -> None:
  36. if have_greenlet:
  37. self.runner = _Runner()
  38. def run(
  39. self,
  40. fn: Callable[..., Coroutine[Any, Any, _T]],
  41. *args: Any,
  42. **kwargs: Any,
  43. ) -> _T:
  44. """Run coroutine on the loop"""
  45. return self.runner.run(fn(*args, **kwargs))
  46. def run_in_greenlet(
  47. self, fn: Callable[..., _T], *args: Any, **kwargs: Any
  48. ) -> _T:
  49. """Run sync function in greenlet. Support nested calls"""
  50. if have_greenlet:
  51. if self.runner.get_loop().is_running():
  52. return fn(*args, **kwargs)
  53. else:
  54. return self.runner.run(greenlet_spawn(fn, *args, **kwargs))
  55. else:
  56. return fn(*args, **kwargs)
  57. def close(self) -> None:
  58. if have_greenlet:
  59. self.runner.close()
  60. if not typing.TYPE_CHECKING and not have_greenlet:
  61. def _not_implemented():
  62. # this conditional is to prevent pylance from considering
  63. # greenlet_spawn() etc as "no return" and dimming out code below it
  64. if have_greenlet:
  65. return None
  66. raise ValueError(
  67. "the greenlet library is required to use this function."
  68. " %s" % greenlet_error
  69. if greenlet_error
  70. else ""
  71. )
  72. def is_exit_exception(e): # noqa: F811
  73. return not isinstance(e, Exception)
  74. def await_only(thing): # type: ignore # noqa: F811
  75. _not_implemented()
  76. def await_fallback(thing): # type: ignore # noqa: F811
  77. return thing
  78. def in_greenlet(): # type: ignore # noqa: F811
  79. _not_implemented()
  80. def greenlet_spawn(fn, *args, **kw): # type: ignore # noqa: F811
  81. _not_implemented()
  82. def AsyncAdaptedLock(*args, **kw): # type: ignore # noqa: F811
  83. _not_implemented()
  84. def _util_async_run(fn, *arg, **kw): # type: ignore # noqa: F811
  85. return fn(*arg, **kw)
  86. def _util_async_run_coroutine_function(fn, *arg, **kw): # type: ignore # noqa: F811,E501
  87. _not_implemented()