_waiter.py 6.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200
  1. # Copyright (c) Microsoft Corporation.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import asyncio
  15. import inspect
  16. import math
  17. import uuid
  18. from asyncio.tasks import Task
  19. from typing import Any, Callable, List, Optional, Tuple, Union
  20. from pyee import EventEmitter
  21. from playwright._impl._connection import ChannelOwner
  22. from playwright._impl._errors import Error, TimeoutError
  23. class Waiter:
  24. def __init__(self, channel_owner: ChannelOwner, event: str) -> None:
  25. self._result: asyncio.Future = asyncio.Future()
  26. self._wait_id = uuid.uuid4().hex
  27. self._loop = channel_owner._loop
  28. self._pending_tasks: List[Task] = []
  29. self._channel = channel_owner._channel
  30. self._registered_listeners: List[Tuple[EventEmitter, str, Callable]] = []
  31. self._logs: List[str] = []
  32. self._wait_for_event_info_before(self._wait_id, event)
  33. def _wait_for_event_info_before(self, wait_id: str, event: str) -> None:
  34. self._channel.send_no_reply(
  35. "waitForEventInfo",
  36. None,
  37. {
  38. "info": {
  39. "waitId": wait_id,
  40. "phase": "before",
  41. "event": event,
  42. }
  43. },
  44. )
  45. def _wait_for_event_info_after(self, wait_id: str, error: Exception = None) -> None:
  46. self._channel._connection.wrap_api_call_sync(
  47. lambda: self._channel.send_no_reply(
  48. "waitForEventInfo",
  49. None,
  50. {
  51. "info": {
  52. "waitId": wait_id,
  53. "phase": "after",
  54. **({"error": str(error)} if error else {}),
  55. },
  56. },
  57. ),
  58. True,
  59. )
  60. def reject_on_event(
  61. self,
  62. emitter: EventEmitter,
  63. event: str,
  64. error: Union[Error, Callable[..., Error]],
  65. predicate: Callable = None,
  66. ) -> None:
  67. def on_match() -> None:
  68. self._reject(error() if callable(error) else error)
  69. def listener(event_data: Any = None) -> None:
  70. self._evaluate_predicate(predicate, event_data, on_match)
  71. emitter.on(event, listener)
  72. self._registered_listeners.append((emitter, event, listener))
  73. def reject_on_timeout(self, timeout: float, message: str) -> None:
  74. if timeout == 0:
  75. return
  76. async def reject() -> None:
  77. await asyncio.sleep(timeout / 1000)
  78. self._reject(TimeoutError(message))
  79. self._pending_tasks.append(self._loop.create_task(reject()))
  80. def _cleanup(self) -> None:
  81. for task in self._pending_tasks:
  82. if not task.done():
  83. task.cancel()
  84. for listener in self._registered_listeners:
  85. listener[0].remove_listener(listener[1], listener[2])
  86. def _fulfill(self, result: Any) -> None:
  87. self._cleanup()
  88. if not self._result.done():
  89. self._result.set_result(result)
  90. self._wait_for_event_info_after(self._wait_id)
  91. def _reject(self, exception: Exception) -> None:
  92. self._cleanup()
  93. if exception:
  94. base_class = TimeoutError if isinstance(exception, TimeoutError) else Error
  95. exception = base_class(str(exception) + format_log_recording(self._logs))
  96. if not self._result.done():
  97. self._result.set_exception(exception)
  98. self._wait_for_event_info_after(self._wait_id, exception)
  99. def wait_for_event(
  100. self,
  101. emitter: EventEmitter,
  102. event: str,
  103. predicate: Callable = None,
  104. ) -> None:
  105. def listener(event_data: Any = None) -> None:
  106. self._evaluate_predicate(
  107. predicate, event_data, lambda: self._fulfill(event_data)
  108. )
  109. emitter.on(event, listener)
  110. self._registered_listeners.append((emitter, event, listener))
  111. def _evaluate_predicate(
  112. self,
  113. predicate: Optional[Callable],
  114. event_data: Any,
  115. on_match: Callable[[], None],
  116. ) -> None:
  117. if predicate is None:
  118. on_match()
  119. return
  120. try:
  121. result = predicate(event_data)
  122. except Exception as e:
  123. self._reject(e)
  124. return
  125. if inspect.iscoroutine(result):
  126. async def _await_predicate(coro: Any) -> None:
  127. try:
  128. matched = await coro
  129. except Exception as e:
  130. self._reject(e)
  131. return
  132. if matched and not self._result.done():
  133. on_match()
  134. self._pending_tasks.append(self._loop.create_task(_await_predicate(result)))
  135. return
  136. if result:
  137. on_match()
  138. def result(self) -> asyncio.Future:
  139. return self._result
  140. def log(self, message: str) -> None:
  141. self._logs.append(message)
  142. try:
  143. self._channel._connection.wrap_api_call_sync(
  144. lambda: self._channel.send_no_reply(
  145. "waitForEventInfo",
  146. None,
  147. {
  148. "info": {
  149. "waitId": self._wait_id,
  150. "phase": "log",
  151. "message": message,
  152. },
  153. },
  154. ),
  155. True,
  156. )
  157. except Exception:
  158. pass
  159. def throw_on_timeout(timeout: float, exception: Exception) -> asyncio.Task:
  160. async def throw() -> None:
  161. await asyncio.sleep(timeout / 1000)
  162. raise exception
  163. return asyncio.create_task(throw())
  164. def format_log_recording(log: List[str]) -> str:
  165. if not log:
  166. return ""
  167. header = " logs "
  168. header_length = 60
  169. left_length = math.floor((header_length - len(header)) / 2)
  170. right_length = header_length - len(header) - left_length
  171. new_line = "\n"
  172. return f"{new_line}{'=' * left_length}{header}{'=' * right_length}{new_line}{new_line.join(log)}{new_line}{'=' * header_length}"