_disposable.py 2.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293
  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 traceback
  17. from typing import Awaitable, Callable, Dict
  18. import greenlet
  19. from playwright._impl._connection import ChannelOwner
  20. from playwright._impl._errors import Error, is_target_closed_error
  21. class Disposable(ChannelOwner):
  22. def __init__(
  23. self, parent: ChannelOwner, type: str, guid: str, initializer: Dict
  24. ) -> None:
  25. super().__init__(parent, type, guid, initializer)
  26. async def dispose(self) -> None:
  27. try:
  28. await self._channel.send(
  29. "dispose",
  30. None,
  31. )
  32. except Exception as e:
  33. if not is_target_closed_error(e):
  34. raise e
  35. async def close(self) -> None:
  36. await self.dispose()
  37. def __repr__(self) -> str:
  38. return "<Disposable>"
  39. class DisposableStub:
  40. def __init__(
  41. self,
  42. dispose_fn: Callable[[], Awaitable[None]],
  43. parent: ChannelOwner,
  44. ) -> None:
  45. self._dispose_fn = dispose_fn
  46. self._loop = parent._loop
  47. self._dispatcher_fiber = parent._dispatcher_fiber
  48. async def dispose(self) -> None:
  49. await self._dispose_fn()
  50. async def __aenter__(self) -> "DisposableStub":
  51. return self
  52. async def __aexit__(self, *args: object) -> None:
  53. await self.dispose()
  54. def __enter__(self) -> "DisposableStub":
  55. return self
  56. def __exit__(self, *args: object) -> None:
  57. self._sync(self.dispose())
  58. def _sync(self, coro: object) -> object:
  59. __tracebackhide__ = True
  60. if self._loop.is_closed():
  61. coro.close() # type: ignore
  62. raise Error("Event loop is closed! Is Playwright already stopped?")
  63. g_self = greenlet.getcurrent()
  64. task = self._loop.create_task(coro) # type: ignore
  65. setattr(task, "__pw_stack__", inspect.stack(0))
  66. setattr(task, "__pw_stack_trace__", traceback.extract_stack(limit=10))
  67. task.add_done_callback(lambda _: g_self.switch())
  68. while not task.done():
  69. self._dispatcher_fiber.switch() # type: ignore
  70. asyncio._set_running_loop(self._loop)
  71. return task.result()
  72. async def close(self) -> None:
  73. await self.dispose()
  74. def __repr__(self) -> str:
  75. return "<Disposable>"