futures.py 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119
  1. # Copyright (C) 2023 The Qt Company Ltd.
  2. # SPDX-License-Identifier: LicenseRef-Qt-Commercial OR LGPL-3.0-only OR GPL-2.0-only OR GPL-3.0-only
  3. from __future__ import annotations
  4. from . import events
  5. from typing import Any, Callable
  6. import asyncio
  7. import contextvars
  8. import enum
  9. class QAsyncioFuture():
  10. """ https://docs.python.org/3/library/asyncio-future.html """
  11. # Declare that this class implements the Future protocol. The field must
  12. # exist and be boolean - True indicates 'await' or 'yield from', False
  13. # indicates 'yield'.
  14. _asyncio_future_blocking = False
  15. class FutureState(enum.Enum):
  16. PENDING = enum.auto()
  17. CANCELLED = enum.auto()
  18. DONE_WITH_RESULT = enum.auto()
  19. DONE_WITH_EXCEPTION = enum.auto()
  20. def __init__(self, *, loop: "events.QAsyncioEventLoop | None" = None,
  21. context: contextvars.Context | None = None) -> None:
  22. self._loop: "events.QAsyncioEventLoop"
  23. if loop is None:
  24. self._loop = asyncio.events.get_event_loop() # type: ignore[assignment]
  25. else:
  26. self._loop = loop
  27. self._context = context
  28. self._state = QAsyncioFuture.FutureState.PENDING
  29. self._result: Any = None
  30. self._exception: BaseException | None = None
  31. self._cancel_message: str | None = None
  32. # List of callbacks that are called when the future is done.
  33. self._callbacks: list[Callable] = list()
  34. def __await__(self):
  35. if not self.done():
  36. self._asyncio_future_blocking = True
  37. yield self
  38. if not self.done():
  39. raise RuntimeError("await was not used with a Future or Future-like object")
  40. return self.result()
  41. __iter__ = __await__
  42. def _schedule_callbacks(self, context: contextvars.Context | None = None):
  43. """ A future can optionally have callbacks that are called when the future is done. """
  44. for cb in self._callbacks:
  45. self._loop.call_soon(
  46. cb, self, context=context if context else self._context)
  47. def result(self) -> Any | Exception:
  48. if self._state == QAsyncioFuture.FutureState.DONE_WITH_RESULT:
  49. return self._result
  50. if self._state == QAsyncioFuture.FutureState.DONE_WITH_EXCEPTION and self._exception:
  51. raise self._exception
  52. if self._state == QAsyncioFuture.FutureState.CANCELLED:
  53. if self._cancel_message:
  54. raise asyncio.CancelledError(self._cancel_message)
  55. else:
  56. raise asyncio.CancelledError
  57. raise asyncio.InvalidStateError
  58. def set_result(self, result: Any) -> None:
  59. self._result = result
  60. self._state = QAsyncioFuture.FutureState.DONE_WITH_RESULT
  61. self._schedule_callbacks()
  62. def set_exception(self, exception: Exception) -> None:
  63. self._exception = exception
  64. self._state = QAsyncioFuture.FutureState.DONE_WITH_EXCEPTION
  65. self._schedule_callbacks()
  66. def done(self) -> bool:
  67. return self._state != QAsyncioFuture.FutureState.PENDING
  68. def cancelled(self) -> bool:
  69. return self._state == QAsyncioFuture.FutureState.CANCELLED
  70. def add_done_callback(self, cb: Callable, *,
  71. context: contextvars.Context | None = None) -> None:
  72. if self.done():
  73. self._loop.call_soon(
  74. cb, self, context=context if context else self._context)
  75. else:
  76. self._callbacks.append(cb)
  77. def remove_done_callback(self, cb: Callable) -> int:
  78. original_len = len(self._callbacks)
  79. self._callbacks = [_cb for _cb in self._callbacks if _cb != cb]
  80. return original_len - len(self._callbacks)
  81. def cancel(self, msg: str | None = None) -> bool:
  82. if self.done():
  83. return False
  84. self._state = QAsyncioFuture.FutureState.CANCELLED
  85. self._cancel_message = msg
  86. self._schedule_callbacks()
  87. return True
  88. def exception(self) -> BaseException | None:
  89. if self._state == QAsyncioFuture.FutureState.CANCELLED:
  90. raise asyncio.CancelledError
  91. if self.done():
  92. return self._exception
  93. raise asyncio.InvalidStateError
  94. def get_loop(self) -> asyncio.AbstractEventLoop:
  95. return self._loop