wait.py 1.4 KB

1234567891011121314151617181920212223242526272829303132333435363738394041
  1. # SPDX-License-Identifier: MIT
  2. from .selectors import (
  3. HAS_SELECT,
  4. DefaultSelector,
  5. EVENT_READ,
  6. EVENT_WRITE
  7. )
  8. def _wait_for_io_events(socks, events, timeout=None):
  9. """ Waits for IO events to be available from a list of sockets
  10. or optionally a single socket if passed in. Returns a list of
  11. sockets that can be interacted with immediately. """
  12. if not HAS_SELECT:
  13. raise ValueError('Platform does not have a selector')
  14. if not isinstance(socks, list):
  15. # Probably just a single socket.
  16. if hasattr(socks, "fileno"):
  17. socks = [socks]
  18. # Otherwise it might be a non-list iterable.
  19. else:
  20. socks = list(socks)
  21. with DefaultSelector() as selector:
  22. for sock in socks:
  23. selector.register(sock, events)
  24. return [key[0].fileobj for key in
  25. selector.select(timeout) if key[1] & events]
  26. def wait_for_read(socks, timeout=None):
  27. """ Waits for reading to be available from a list of sockets
  28. or optionally a single socket if passed in. Returns a list of
  29. sockets that can be read from immediately. """
  30. return _wait_for_io_events(socks, EVENT_READ, timeout)
  31. def wait_for_write(socks, timeout=None):
  32. """ Waits for writing to be available from a list of sockets
  33. or optionally a single socket if passed in. Returns a list of
  34. sockets that can be written to immediately. """
  35. return _wait_for_io_events(socks, EVENT_WRITE, timeout)