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