async def read_stream(eos: asyncio.Future,
rdr: Coroutine) -> Union[_EndOfStream, Any]:
try:
done, _pending = await asyncio.wait(
{eos, asyncio.create_task(rdr)},
return_when=asyncio.FIRST_COMPLETED)
if len(done) <= 1 and eos in done:
# done without unprocessed item
await eos # reraise exception if that caused eos
return EndOfStream
for fut in done:
if fut is not eos:
if fut.cancelled():
if eos.done():
return EndOfStream
return await fut
except asyncio.CancelledError:
if eos.done():
await eos # reraise exception if that caused eos
return EndOfStream
logger.debug(
f"stream reading cancelled before eos: {eos!s}",
exc_info=True,
stack_info=True,
)
raise # reraise as is, if not at eos yet
assert False, "impossible to reach here"
def run_until_complete(self, future: asyncio.Future) -> Any:
self._check_running()
new_task = False
if not asyncio.isfuture(future):
future = asyncio.ensure_future(future, loop=self)
# We wrapped `future` in a new Task since it was not a Future.
new_task = True
# An exception is raised if the new task doesn't complete, so there is no need to log the "destroy
# pending task" message.
future._log_destroy_pending = False
else:
if _helpers.get_future_loop(future) is not self:
raise ValueError("Future does not belong to this loop")
future.add_done_callback(_run_until_complete_cb)
try:
self._run_mainloop()
except:
if new_task and future.done() and not future.cancelled():
# The coroutine raised a BaseException. Consume the exception to not log a warning (Future
# will log a warning if its exception is not retrieved), the caller doesn't have access to the
# task wrapper we made.
future.exception()
raise
finally:
future.remove_done_callback(_run_until_complete_cb)
if not future.done():
raise RuntimeError('Event loop stopped before Future completed.')
return future.result()
class Stream:
def __init__(self):
self._future = Future()
self._complete = Event()
self._listeners = [ ]
def __aiter__(self):
return self
async def __anext__(self):
if self._complete.is_set():
raise StopAsyncIteration
result = await self._future
self._future = Future()
return result
def write(self, item, last):
if self._complete.is_set():
return
if last:
self._set_complete()
if self._future.done():
self._future = Future()
self._future.set_result(item)
def abort(self, exc):
if self._complete.is_set():
return
self._set_complete()
if self._future.done():
self._future = Future()
self._future.set_exception(exc)
def cancel(self):
self._set_complete()
if self._future.done():
self._future = Future()
self._future.cancel()
async def completed(self):
await self._complete.wait()
@property
def is_complete(self):
return self._complete.is_set()
def add_complete_listener(self, listener):
self._listeners.append(listener)
def _set_complete(self):
self._complete.set()
for listener in self._listeners:
listener()
def _in_thread(future: asyncio.Future, func: Callable,
loop: asyncio.AbstractEventLoop):
try:
result = func()
except Exception as e:
if future.done():
return
loop.call_soon_threadsafe(future.set_exception, e)
else:
if future.done():
return
loop.call_soon_threadsafe(future.set_result, result)
def copy_result(from_: Future, to: Future) -> None:
if not from_.done():
raise ValueError('from_ must be a completed Future')
if to.done():
raise ValueError('to must NOT be a completed Future')
if from_.cancelled():
to.cancel()
else:
exception = from_.exception()
if exception is not None:
to.set_exception(exception)
else:
result = from_.result()
to.set_result(result)
def set_future_result(future: Future,
result: Any) -> None: # type: ignore # pypy
if not future.done():
if isinstance(result, Exception):
future.set_exception(result)
else:
future.set_result(result)
async def test_handle_response_none(self):
transport = FAsyncTransport(1024)
ctx = FContext()
future = Future()
transport._futures[str(ctx._get_op_id())] = future
await transport.handle_response(None)
self.assertFalse(future.done())
class Tester:
def __init__(self):
self.end = Future()
def __call__(self, *args, **kwargs):
if not self.end.done():
self.end.set_result((args, kwargs))
def _on_child_task_done(self, parent_task: asyncio.Future,
child_task: asyncio.Future):
parent_task.child_count -= 1 # type: ignore
try:
metadata = child_task.result()
if metadata and not parent_task.done():
parent_task.set_result(metadata)
except asyncio.CancelledError:
pass
except Exception as e:
logging.debug("Child task resulted in exception")
logger.exception(e)
if parent_task.child_count <= 0 and not parent_task.done(
): # type: ignore
parent_task.set_result(None)
async def _send_udp_message_with_retry(
message: str,
transport: asyncio.DatagramTransport,
response_future: asyncio.Future,
ip: str,
port: int,
) -> None:
"""Send a UDP message multiple times until we reach the maximum or a response is recieved."""
data = message.encode("utf-8")
send_wait = FIRST_SEND_INTERVAL
total_waited = 0.0
for send in range(MAX_SEND_DATAGRAMS):
if transport.is_closing() or response_future.done():
return
attempt = send + 1
_LOGGER.debug(
"%s: >> %s (%s/%s) backoff=%s",
ip,
data,
attempt,
MAX_SEND_DATAGRAMS,
total_waited,
)
transport.sendto(data, (ip, port))
await asyncio.sleep(send_wait)
total_waited += send_wait
send_wait = min(send_wait * 2, MAX_BACKOFF)
class AIterator:
def __init__(self):
self.notifications = []
self.future = Future()
source.pipe(ops.materialize()).subscribe(self.on_next)
def feeder(self):
if not self.notifications or self.future.done():
return
notification = self.notifications.pop(0)
dispatch = {
'N':
lambda: self.future.set_result(notification.value),
'E':
lambda: self.future.set_exception(notification.
exception),
'C':
lambda: self.future.set_exception(
StopAsyncIteration)
}
dispatch[notification.kind]()
def on_next(self, notification):
self.notifications.append(notification)
self.feeder()
async def __anext__(self):
self.feeder()
value = await self.future
self.future = Future()
return value
def _run_coro(future: asyncio.Future, coro: Awaitable[Any],
callback: Callable[[str, bool], Awaitable[Any]]) -> None:
loop = asyncio.new_event_loop()
task = loop.create_task(coro)
userge.loop.call_soon_threadsafe(
future.add_done_callback, lambda _: (loop.is_running(
) and future.cancelled() and loop.call_soon_threadsafe(task.cancel)))
try:
ret, exc = None, None
with redirect() as out:
try:
ret = loop.run_until_complete(task)
except asyncio.CancelledError:
return
except Exception: # pylint: disable=broad-except
exc = traceback.format_exc().strip()
output = exc or out.getvalue()
if ret is not None:
output += str(ret)
loop.run_until_complete(callback(output, exc is not None))
finally:
loop.run_until_complete(loop.shutdown_asyncgens())
loop.close()
userge.loop.call_soon_threadsafe(
lambda: future.done() or future.set_result(None))
def demo():
f = Future()
print(f.done())
async def main(f: asyncio.Future):
await asyncio.sleep(1)
f.set_result('I have finished.')
return 13
loop = asyncio.get_event_loop()
fut = asyncio.Future()
# Schedule the main() coroutine, passing the future
task = loop.create_task(main(fut))
print(fut.done())
print(task.done())
# Here we use run_until_complete() on a Future instance, rather than a Task instance
loop.run_until_complete(fut)
print(fut.done())
print(task.done())
print(fut.result())
print(task.result())
def _on_result(future: asyncio.Future, new_future: asyncio.Future = None):
if not new_future.done():
exc = future.exception()
if exc:
return new_future.set_exception(exc)
new_future.set_result(future.result())
async def get(self, timeout=None) -> Connection:
async with self._alock:
if len(self._idle_pool) > 0:
c = self._idle_pool.pop()
self._using_pool.add(c)
_connectionCtxVar.set(c)
return c
if len(self._using_pool) < self._pool_size:
c = await self._aopen()
self._using_pool.add(c)
_connectionCtxVar.set(c)
return c
future = Future()
self._waiting.append(future)
if timeout and timeout > 0:
asyncio.get_event_loop().call_later(timeout, lambda: future.done())
await future
c = future.result()
if not c:
return c
_connectionCtxVar.set(c)
self._using_pool.add(c)
return c
class MockTransport(MockObject):
"""
This object mocks the transport of the socks5 connection.
"""
def __init__(self, loop):
self.connected = True
self.written_data = []
self.host = '123.123.123.123'
self.ip = 123
self.num_messages = 2
self.done = Future(loop=loop)
def close(self):
self.connected = False
def write(self, data):
self.written_data.append(data)
if len(self.written_data
) == self.num_messages and not self.done.done():
self.done.set_result(None)
def get_extra_info(self, *_):
return self.host, self.ip
async def wait_until_done(self, timeout=1):
await wait_for(self.done, timeout=timeout)
async def consume_messages(self, when_connected: asyncio.Future,
subscriptions: Subscriptions):
"""Consumes messages form the open web socket connection and forwards them to the callback"""
self.connection = await websockets.connect(self.host, ssl=True)
if subscriptions is None:
# When no subscriptions are added we complete on connect
when_connected.set_result("Connected")
if self.api_token is not None:
auth_message = self.get_auth_message_as_json()
await self.send_message(auth_message)
else:
# Subscribe immediately if no api_token is provided - public channels need no token
await self.subscribe(subscriptions)
async for raw_message in self.connection:
LOG.debug(">>> %s", raw_message)
json_message = json.loads(raw_message)
# Subscribe only after authentication confirmation
if json_message["type"] == "AUTHENTICATED":
LOG.debug("Successfully authenticated! Subscribe now.")
await self.subscribe(subscriptions)
elif when_connected.done() is False and json_message["type"] == "SUBSCRIPTIONS" \
and subscriptions is not None:
# Complete start on subscription confirmation message
when_connected.set_result("Connected with subscriptions")
elif json_message["type"] == "ERROR":
LOG.error("!!! %s", json_message)
# Forward all events
await self.callback(json_message)
def _stop_task(key: str, future: asyncio.Future) -> None:
if not future.cancelled():
future.cancel()
logging.info(f"stop task:{key}")
elif future.done():
logging.warning(f"task:{key} already stop")
else:
logging.warning(f"{key} can't stop")
class ProgressStream:
def __init__(self):
self._progress = Queue()
self._progress_task = create_task(self._progress.get())
self._complete = Future()
self._complete_task = create_task(self._complete)
def __aiter__(self):
return self
async def __anext__(self):
if self._complete.done():
raise StopAsyncIteration
done, pending = await wait({self._complete_task, self._progress_task},
return_when=FIRST_COMPLETED)
if self._complete_task in done:
self._progress_task.cancel()
self._complete_task.result(
) # throws an exception (when necessary)
raise StopAsyncIteration
else:
progress = self._progress_task.result()
self._progress_task = create_task(self._progress.get())
return progress
def write(self, item):
if self._complete.done():
raise InvalidStateError
if self._progress.qsize() > 0:
self._progress.get_nowait()
self._progress.put_nowait(item)
def set_result(self, result):
self._complete.set_result(result)
def set_exception(self, e):
self._complete.set_exception(e)
def result(self):
return self._complete.result()
def done(self):
return self._complete.done()
def set_result(future: Future, result: Any, exception: Exception) -> None:
if future.done():
return
if exception:
future.set_exception(exception)
else:
future.set_result(result)
class Tester:
def __init__(self):
self.end = Future()
def __call__(self, *args, **kwargs):
if not self.end.done():
self.end.set_result((args, kwargs))
def propagate(from_: Future, to: Future) -> None:
"""
Copy the value that ``from_`` is completed with to ``to``, whenever ``from_`` is completed.
"""
if from_.done():
copy_result(from_, to)
else:
from_.add_done_callback(lambda _: copy_result(from_, to))
async def test_handle_response_unregistered_op_id(self):
transport = FAsyncTransport(1024)
ctx1 = FContext()
ctx2 = FContext()
future = Future()
transport._futures[str(ctx1._get_op_id())] = future
await transport.handle_response(utils.mock_frame(ctx2))
self.assertFalse(future.done())
async def _awaiter(future: Future) -> T:
try:
result = await future
return result
except asyncio.CancelledError as e:
if not future.done():
future.set_exception(e)
raise
async def awaiter(future: asyncio.Future) -> T:
"""等待future执行结束，返回执行结果 ."""
try:
result = await future
return result
except asyncio.CancelledError as e:
if not future.done():
future.set_exception(e)
raise
def _create_job_callback(job: Job, fut: asyncio.Future):
exc = fut.exception()
if not fut.cancelled() and exc:
raise exc
if fut.done():
job.meta_update(result=fut.result())
job.save()
def callback(my_future: Future):
if not my_future.done() or my_future.cancelled():
return
result: ExResult = my_future.result()
if result.solver_result.satisfiable:
return
else:
print("kept graph")
with open("results_random_planar_110_filtered.json", mode="a") as f:
print(simplejson.dumps(result), file=f)
def _finalise_future(fut: Future) -> Optional[Union[Response, str]]:
if fut.done():
err = fut.exception()
if err:
error_logger.error("%s", err, exc_info=err)
else:
return fut.result()
else:
fut.cancel()
async def _pipe_to_companion(
in_stream: Stream[_TSend, _TRecv],
out_stream: Stream[_TRecv, _TSend],
started_future: asyncio.Future,
) -> None:
async for message in in_stream:
await out_stream.send_message(message)
if not started_future.done():
started_future.set_result(None)
await out_stream.end()
class MySubHandler:
"""
More advanced subscription client using Future, so we can await events in tests.
"""
def __init__(self):
self.future = Future()
def reset(self):
self.future = Future()
async def result(self):
return await wait_for(self.future, 2)
def datachange_notification(self, node, val, data):
if not self.future.done():
self.future.set_result((node, val, data))
def event_notification(self, event):
if not self.future.done():
self.future.set_result(event)
def run_job_done_callback(f: asyncio.Future):
if f.cancelled():
if not job.done():
job.cancel()
elif f.exception():
job.set_exception(f.exception())
elif f.done():
job.set_result(f.result())
# cancel should be triggered by job
else:
raise RuntimeError(f"Unexpected future {f}")
def _send_request(self, app_name, endpoint, entity, params, timeout):
packet = MessagePacket.request(self.name, self.version, app_name, _Service._REQ_PKT_STR, endpoint, params,
entity)
future = Future()
request_id = params['request_id']
self._pending_requests[request_id] = future
try:
self.tcp_bus.send(packet)
except ClientException:
if not future.done() and not future.cancelled():
error = 'Client not found'
exception = ClientException(error)
exception.error = error
future.set_exception(exception)
_Service.time_future(future, timeout)
return future
def _send_request(self, app_name, endpoint, entity, params):
packet = MessagePacket.request(self.name, self.version, app_name, _Service._REQ_PKT_STR, endpoint, params,
entity)
future = Future()
request_id = params['request_id']
self._pending_requests[request_id] = future
try:
self.tcp_bus.send(packet)
except ClientException as e:
if not future.done() and not future.cancelled():
ERROR = '101_Client not found'
exception = RequestException(ERROR)
exception.error = ERROR
future.set_exception(exception)
_Service.time_future(future, TCPServiceClient.REQUEST_TIMEOUT_SECS)
return future
class SubscriberThread(threading.Thread):
def __init__(self, topic, id):
super().__init__(target=self)
self.subscriber = Subscriber(topic=topic, id=id)
self.future = Future()
def done(self):
return self.future.done()
def result(self):
return self.future.result()
def run(self):
try:
self.subscriber.open()
result = self.loop()
self.future.set_result(result)
except Exception as error:
self.future.set_exception(error)
finally:
self.subscriber.close()
def loop(self):
raise NotImplementedError()
class WebSocketAdapterProtocol(asyncio.Protocol):
"""
Adapter class for asyncio-based WebSocket client and server protocols.
"""
def connection_made(self, transport):
self.transport = transport
self.receive_queue = deque()
self._consume()
try:
self.peer = peer2str(transport.get_extra_info('peername'))
except:
self.peer = u"?"
self._connectionMade()
def connection_lost(self, exc):
self._connectionLost(exc)
# according to asyncio docs, connection_lost(None) is called
# if something else called transport.close()
if exc is not None:
self.transport.close()
self.transport = None
def _consume(self):
self.waiter = Future(loop=self.factory.loop or txaio.config.loop)
def process(_):
while len(self.receive_queue):
data = self.receive_queue.popleft()
if self.transport:
self._dataReceived(data)
self._consume()
self.waiter.add_done_callback(process)
def data_received(self, data):
self.receive_queue.append(data)
if not self.waiter.done():
self.waiter.set_result(None)
def _closeConnection(self, abort=False):
if abort and hasattr(self.transport, 'abort'):
self.transport.abort()
else:
self.transport.close()
def _onOpen(self):
res = self.onOpen()
if yields(res):
ensure_future(res)
def _onMessageBegin(self, isBinary):
res = self.onMessageBegin(isBinary)
if yields(res):
ensure_future(res)
def _onMessageFrameBegin(self, length):
res = self.onMessageFrameBegin(length)
if yields(res):
ensure_future(res)
def _onMessageFrameData(self, payload):
res = self.onMessageFrameData(payload)
if yields(res):
ensure_future(res)
def _onMessageFrameEnd(self):
res = self.onMessageFrameEnd()
if yields(res):
ensure_future(res)
def _onMessageFrame(self, payload):
res = self.onMessageFrame(payload)
if yields(res):
ensure_future(res)
def _onMessageEnd(self):
res = self.onMessageEnd()
if yields(res):
ensure_future(res)
def _onMessage(self, payload, isBinary):
res = self.onMessage(payload, isBinary)
if yields(res):
ensure_future(res)
def _onPing(self, payload):
res = self.onPing(payload)
if yields(res):
ensure_future(res)
def _onPong(self, payload):
res = self.onPong(payload)
if yields(res):
ensure_future(res)
def _onClose(self, wasClean, code, reason):
res = self.onClose(wasClean, code, reason)
if yields(res):
ensure_future(res)
def registerProducer(self, producer, streaming):
raise Exception("not implemented")
def unregisterProducer(self):
# note that generic websocket/protocol.py code calls
# .unregisterProducer whenever we dropConnection -- that's
# correct behavior on Twisted so either we'd have to
# try/except there, or special-case Twisted, ..or just make
# this "not an error"
pass
class WebSocketAdapterProtocol(asyncio.Protocol):
"""
Adapter class for asyncio-based WebSocket client and server protocols.
"""
def connection_made(self, transport):
self.transport = transport
self.receive_queue = deque()
self._consume()
try:
peer = transport.get_extra_info('peername')
try:
# FIXME: tcp4 vs tcp6
self.peer = u"tcp:%s:%d" % (peer[0], peer[1])
except:
# e.g. Unix Domain sockets don't have host/port
self.peer = u"unix:{0}".format(peer)
except:
self.peer = u"?"
self._connectionMade()
def connection_lost(self, exc):
self._connectionLost(exc)
# according to asyncio docs, connection_lost(None) is called
# if something else called transport.close()
if exc is not None:
self.transport.close()
self.transport = None
def _consume(self):
self.waiter = Future(loop=self.factory.loop or txaio.config.loop)
def process(_):
while len(self.receive_queue):
data = self.receive_queue.popleft()
if self.transport:
self._dataReceived(data)
self._consume()
self.waiter.add_done_callback(process)
def data_received(self, data):
self.receive_queue.append(data)
if not self.waiter.done():
self.waiter.set_result(None)
# noinspection PyUnusedLocal
def _closeConnection(self, abort=False):
self.transport.close()
def _onOpen(self):
res = self.onOpen()
if yields(res):
ensure_future(res)
def _onMessageBegin(self, isBinary):
res = self.onMessageBegin(isBinary)
if yields(res):
ensure_future(res)
def _onMessageFrameBegin(self, length):
res = self.onMessageFrameBegin(length)
if yields(res):
ensure_future(res)
def _onMessageFrameData(self, payload):
res = self.onMessageFrameData(payload)
if yields(res):
ensure_future(res)
def _onMessageFrameEnd(self):
res = self.onMessageFrameEnd()
if yields(res):
ensure_future(res)
def _onMessageFrame(self, payload):
res = self.onMessageFrame(payload)
if yields(res):
ensure_future(res)
def _onMessageEnd(self):
res = self.onMessageEnd()
if yields(res):
ensure_future(res)
def _onMessage(self, payload, isBinary):
res = self.onMessage(payload, isBinary)
if yields(res):
ensure_future(res)
def _onPing(self, payload):
res = self.onPing(payload)
if yields(res):
ensure_future(res)
def _onPong(self, payload):
res = self.onPong(payload)
if yields(res):
ensure_future(res)
def _onClose(self, wasClean, code, reason):
res = self.onClose(wasClean, code, reason)
if yields(res):
ensure_future(res)
def get_channel_id(self):
"""
Implements :func:`autobahn.wamp.interfaces.ITransport.get_channel_id`
"""
self.log.debug('FIXME: transport channel binding not implemented for asyncio (autobahn-python issue #729)')
return None
def registerProducer(self, producer, streaming):
raise Exception("not implemented")
class WebSocketAdapterProtocol(asyncio.Protocol):
"""
Adapter class for asyncio-based WebSocket client and server protocols.
"""
def connection_made(self, transport):
self.transport = transport
self.receive_queue = deque()
self._consume()
try:
peer = transport.get_extra_info('peername')
try:
# FIXME: tcp4 vs tcp6
self.peer = "tcp:%s:%d" % (peer[0], peer[1])
except:
# e.g. Unix Domain sockets don't have host/port
self.peer = "unix:{0}".format(peer)
except:
self.peer = "?"
self._connectionMade()
def connection_lost(self, exc):
self._connectionLost(exc)
self.transport = None
def _consume(self):
self.waiter = Future()
def process(_):
while len(self.receive_queue):
data = self.receive_queue.popleft()
if self.transport:
self._dataReceived(data)
else:
print("WebSocketAdapterProtocol._consume: no transport")
self._consume()
self.waiter.add_done_callback(process)
def data_received(self, data):
self.receive_queue.append(data)
if not self.waiter.done():
self.waiter.set_result(None)
# noinspection PyUnusedLocal
def _closeConnection(self, abort=False):
self.transport.close()
def _onOpen(self):
res = self.onOpen()
if yields(res):
asyncio.async(res)
def _onMessageBegin(self, isBinary):
res = self.onMessageBegin(isBinary)
if yields(res):
asyncio.async(res)
def _onMessageFrameBegin(self, length):
res = self.onMessageFrameBegin(length)
if yields(res):
asyncio.async(res)
def _onMessageFrameData(self, payload):
res = self.onMessageFrameData(payload)
if yields(res):
asyncio.async(res)
def _onMessageFrameEnd(self):
res = self.onMessageFrameEnd()
if yields(res):
asyncio.async(res)
def _onMessageFrame(self, payload):
res = self.onMessageFrame(payload)
if yields(res):
asyncio.async(res)
def _onMessageEnd(self):
res = self.onMessageEnd()
if yields(res):
asyncio.async(res)
def _onMessage(self, payload, isBinary):
res = self.onMessage(payload, isBinary)
if yields(res):
asyncio.async(res)
def _onPing(self, payload):
res = self.onPing(payload)
if yields(res):
asyncio.async(res)
def _onPong(self, payload):
res = self.onPong(payload)
if yields(res):
asyncio.async(res)
def _onClose(self, wasClean, code, reason):
res = self.onClose(wasClean, code, reason)
if yields(res):
asyncio.async(res)
def registerProducer(self, producer, streaming):
raise Exception("not implemented")
