_compat_py3k.py 2.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667
  1. # util/_compat_py3k.py
  2. # Copyright (C) 2005-2022 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. from functools import wraps
  8. # vendored from py3.7
  9. class _AsyncGeneratorContextManager:
  10. """Helper for @asynccontextmanager."""
  11. def __init__(self, func, args, kwds):
  12. self.gen = func(*args, **kwds)
  13. self.func, self.args, self.kwds = func, args, kwds
  14. doc = getattr(func, "__doc__", None)
  15. if doc is None:
  16. doc = type(self).__doc__
  17. self.__doc__ = doc
  18. async def __aenter__(self):
  19. try:
  20. return await self.gen.__anext__()
  21. except StopAsyncIteration:
  22. raise RuntimeError("generator didn't yield") from None
  23. async def __aexit__(self, typ, value, traceback):
  24. if typ is None:
  25. try:
  26. await self.gen.__anext__()
  27. except StopAsyncIteration:
  28. return
  29. else:
  30. raise RuntimeError("generator didn't stop")
  31. else:
  32. if value is None:
  33. value = typ()
  34. # See _GeneratorContextManager.__exit__ for comments on subtleties
  35. # in this implementation
  36. try:
  37. await self.gen.athrow(typ, value, traceback)
  38. raise RuntimeError("generator didn't stop after athrow()")
  39. except StopAsyncIteration as exc:
  40. return exc is not value
  41. except RuntimeError as exc:
  42. if exc is value:
  43. return False
  44. if isinstance(value, (StopIteration, StopAsyncIteration)):
  45. if exc.__cause__ is value:
  46. return False
  47. raise
  48. except BaseException as exc:
  49. if exc is not value:
  50. raise
  51. # using the vendored version in all cases at the moment to establish
  52. # full test coverage
  53. def asynccontextmanager(func):
  54. @wraps(func)
  55. def helper(*args, **kwds):
  56. return _AsyncGeneratorContextManager(func, args, kwds)
  57. return helper