Python asyncio.ensure_future() Examples

def async_run(tasks):
    """
    run a group of tasks async
    Requires the tasks arg to be a list of functools.partial()
    """
    if not tasks:
        return

    # start a new async event loop
    loop = asyncio.get_event_loop()
    # https://github.com/python/asyncio/issues/258
    executor = concurrent.futures.ThreadPoolExecutor(5)
    loop.set_default_executor(executor)

    async_tasks = [asyncio.ensure_future(async_task(task, loop)) for task in tasks]
    # run tasks in parallel
    loop.run_until_complete(asyncio.wait(async_tasks))
    # deal with errors (exceptions, etc)
    for task in async_tasks:
        error = task.exception()
        if error is not None:
            raise error

    executor.shutdown(wait=True) 

def test_websocket_non_regression_bug_105(event_loop, server):

    # This test will check a fix to a race condition which happens if the user is trying
    # to connect using the same client twice at the same time
    # See bug #105

    url = f"ws://{server.hostname}:{server.port}/graphql"
    print(f"url = {url}")

    sample_transport = WebsocketsTransport(url=url)

    client = Client(transport=sample_transport)

    # Create a coroutine which start the connection with the transport but does nothing
    async def client_connect(client):
        async with client:
            await asyncio.sleep(2 * MS)

    # Create two tasks which will try to connect using the same client (not allowed)
    connect_task1 = asyncio.ensure_future(client_connect(client))
    connect_task2 = asyncio.ensure_future(client_connect(client))

    with pytest.raises(TransportAlreadyConnected):
        await asyncio.gather(connect_task1, connect_task2) 

def main(_):
  urls = get_urls_for_shard_group(
      FLAGS.urls_dir, FLAGS.shard_id, FLAGS.group_id)
  tf.logging.info("Fetching %d URLs for shard %d, group %d",
                  len(urls), FLAGS.shard_id, FLAGS.group_id)

  tf.gfile.MakeDirs(FLAGS.out_dir)
  out_fname = tfrecord_fname(FLAGS.out_dir, FLAGS.shard_id)

  with utils.timing("group_fetch"):
    logging_fnames = {}
    if FLAGS.log_samples:
      logging_fnames["samples"] = os.path.join(
          FLAGS.out_dir, "samples.%d.txt" % FLAGS.shard_id)
    loop = asyncio.get_event_loop()
    num_written = loop.run_until_complete(asyncio.ensure_future(
        fetch_urls(urls,
                   out_fname,
                   logging_fnames)))

  tf.logging.info("Total URLs: %d", len(urls))
  tf.logging.info("Num written: %d", num_written)
  tf.logging.info("Coverage: %.1f", (num_written / len(urls)) * 100) 

def test_task_local() -> None:
    local_ = TaskLocal()
    queue: asyncio.Queue = asyncio.Queue()
    tasks = 2
    for _ in range(tasks):
        queue.put_nowait(None)

    async def _test_local(value: int) -> int:
        local_.test = value
        await queue.get()
        queue.task_done()
        await queue.join()
        return local_.test

    futures = [asyncio.ensure_future(_test_local(value)) for value in range(tasks)]
    asyncio.gather(*futures)
    for value, future in enumerate(futures):
        assert (await future) == value 

def test_copy_current_app_context() -> None:
    app = Quart(__name__)

    @app.route("/")
    async def index() -> str:
        g.foo = "bar"  # type: ignore

        @copy_current_app_context
        async def within_context() -> None:
            assert g.foo == "bar"

        await asyncio.ensure_future(within_context())
        return ""

    test_client = app.test_client()
    response = await test_client.get("/")
    assert response.status_code == 200 

def test_copy_current_websocket_context() -> None:
    app = Quart(__name__)

    @app.websocket("/")
    async def index() -> None:
        @copy_current_websocket_context
        async def within_context() -> None:
            return websocket.path

        data = await asyncio.ensure_future(within_context())
        await websocket.send(data.encode())

    test_client = app.test_client()
    async with test_client.websocket("/") as test_websocket:
        data = await test_websocket.receive()
    assert cast(bytes, data) == b"/" 

def add_task(self, task):
        """Schedule a task to run later, after the loop has started.
        Different from asyncio.ensure_future in that it does not
        also return a future, and the actual ensure_future call
        is delayed until before server start.

        :param task: future, couroutine or awaitable
        """
        try:
            loop = self.loop  # Will raise SanicError if loop is not started
            self._loop_add_task(task, self, loop)
        except SanicException:
            self.listener("before_server_start")(
                partial(self._loop_add_task, task)
            )

    # Decorator 

def stream_complete(self, stream_id: int):
        """
        When a stream is complete, we can send our response.
        """
        try:
            request_data = self.stream_data[stream_id]
        except KeyError:
            # Just return, we probably 405'd this already
            return

        headers = request_data.headers
        body = request_data.data.getvalue().decode('utf-8')

        data = json.dumps(
            {"headers": headers, "body": body}, indent=4
        ).encode("utf8")

        response_headers = (
            (':status', '200'),
            ('content-type', 'application/json'),
            ('content-length', str(len(data))),
            ('server', 'asyncio-h2'),
        )
        self.conn.send_headers(stream_id, response_headers)
        asyncio.ensure_future(self.send_data(data, stream_id)) 

def test_kv_missing(self, loop, consul_port):
        async def main():
            c = consul.aio.Consul(port=consul_port, loop=loop)

            fut = asyncio.ensure_future(put(), loop=loop)
            await c.kv.put('index', 'bump')
            index, data = await c.kv.get('foo')
            assert data is None
            index, data = await c.kv.get('foo', index=index)
            assert data['Value'] == six.b('bar')
            await fut

        async def put():
            c = consul.aio.Consul(port=consul_port, loop=loop)

            await asyncio.sleep(2.0 / 100, loop=loop)
            await c.kv.put('foo', 'bar')

        loop.run_until_complete(main()) 

def test_kv_subscribe(self, loop, consul_port):
        async def get():
            c = consul.aio.Consul(port=consul_port, loop=loop)
            fut = asyncio.ensure_future(put(), loop=loop)
            index, data = await c.kv.get('foo')
            assert data is None
            index, data = await c.kv.get('foo', index=index)
            assert data['Value'] == six.b('bar')
            await fut

        async def put():
            c = consul.aio.Consul(port=consul_port, loop=loop)
            await asyncio.sleep(1.0 / 100, loop=loop)
            response = await c.kv.put('foo', 'bar')
            assert response is True

        loop.run_until_complete(get()) 

def test_session(self, loop, consul_port):
        async def monitor():
            c = consul.aio.Consul(port=consul_port, loop=loop)
            fut = asyncio.ensure_future(register(), loop=loop)
            index, services = await c.session.list()
            assert services == []
            await asyncio.sleep(20 / 1000.0, loop=loop)

            index, services = await c.session.list(index=index)
            assert len(services)

            index, services = await c.session.list(index=index)
            assert services == []
            await fut

        async def register():
            c = consul.aio.Consul(port=consul_port, loop=loop)
            await asyncio.sleep(1.0 / 100, loop=loop)
            session_id = await c.session.create()
            await asyncio.sleep(50 / 1000.0, loop=loop)
            response = await c.session.destroy(session_id)
            assert response is True

        loop.run_until_complete(monitor()) 

def fetch_bar2():
    day = datetime.datetime.strptime('20160114', '%Y%m%d').replace(tzinfo=pytz.FixedOffset(480))
    end = datetime.datetime.strptime('20160914', '%Y%m%d').replace(tzinfo=pytz.FixedOffset(480))
    while day <= end:
        day, trading = await is_trading_day(day)
        if trading:
            print('process ', day)
            tasks = [
                asyncio.ensure_future(update_from_shfe(day)),
                asyncio.ensure_future(update_from_dce(day)),
                asyncio.ensure_future(update_from_czce(day)),
                asyncio.ensure_future(update_from_cffex(day)),
            ]
            await asyncio.wait(tasks)
        day += datetime.timedelta(days=1)
    print('all done!')


# asyncio.get_event_loop().run_until_complete(fetch_bar2())
# create_main_all()
# fetch_from_quandl_all()
# clean_dailybar()
# load_kt_data() 

def test_retry_with_asyncio(mock_client, mock_response):
    import asyncio

    @asyncio.coroutine
    def coroutine():
        return mock_response

    # Setup
    mock_response.with_json({"id": 123, "name": "prkumar"})
    mock_client.with_side_effect([Exception, coroutine()])
    mock_client.with_io(io.AsyncioStrategy())
    github = GitHub(base_url=BASE_URL, client=mock_client)

    # Run
    awaitable = github.get_user("prkumar")
    loop = asyncio.get_event_loop()
    response = loop.run_until_complete(asyncio.ensure_future(awaitable))

    # Verify
    assert len(mock_client.history) == 2
    assert response.json() == {"id": 123, "name": "prkumar"} 

def test_request_send(self, mocker, aiohttp_session_mock):
        # Setup
        import asyncio

        expected_response = mocker.Mock()

        @asyncio.coroutine
        def request(*args, **kwargs):
            return expected_response

        aiohttp_session_mock.request = request
        client = aiohttp_.AiohttpClient(aiohttp_session_mock)

        # Run
        response = client.send((1, 2, {}))
        loop = asyncio.get_event_loop()
        value = loop.run_until_complete(asyncio.ensure_future(response))

        # Verify
        assert value == expected_response 

def test_create(self, mocker):
        # Setup
        import asyncio

        session_cls_mock = mocker.patch("aiohttp.ClientSession")
        positionals = [1]
        keywords = {"keyword": 2}

        # Run: Create client
        client = aiohttp_.AiohttpClient.create(*positionals, **keywords)

        # Verify: session hasn't been created yet.
        assert not session_cls_mock.called

        # Run: Get session
        loop = asyncio.get_event_loop()
        loop.run_until_complete(asyncio.ensure_future(client.session()))

        # Verify: session created with args
        session_cls_mock.assert_called_with(*positionals, **keywords) 

def import_file(filename):
    log.info("import_file: {}".format(filename))
    loop = globals["loop"]
    max_concurrent_tasks = config.get("max_concurrent_tasks")
    tasks = []
    with open(filename, 'r') as fh:
        for line in fh:
            line = line.rstrip()
            #loop.run_until_complete(import_line(line))
            tasks.append(asyncio.ensure_future(import_line(line)))
            if len(tasks) < max_concurrent_tasks:
                continue  # get next line
            # got a batch, move them out!
            loop.run_until_complete(asyncio.gather(*tasks))
            tasks = []
    # finish any remaining tasks
    loop.run_until_complete(asyncio.gather(*tasks))
    globals["files_read"] += 1 

def main():
    pub = await aioredis.create_redis(
        'redis://localhost')
    sub = await aioredis.create_redis(
        'redis://localhost')
    res = await sub.subscribe('chan:1')
    ch1 = res[0]

    tsk = asyncio.ensure_future(reader(ch1))

    res = await pub.publish_json('chan:1', ["Hello", "world"])
    assert res == 1

    await sub.unsubscribe('chan:1')
    await tsk
    sub.close()
    pub.close() 

def main():
    loop = asyncio.get_event_loop()
    tsk = asyncio.ensure_future(pubsub(), loop=loop)

    async def publish():
        pub = await aioredis.create_redis(
            'redis://localhost')
        while not tsk.done():
            # wait for clients to subscribe
            while True:
                subs = await pub.pubsub_numsub('channel:1')
                if subs[b'channel:1'] == 1:
                    break
                await asyncio.sleep(0, loop=loop)
            # publish some messages
            for msg in ['one', 'two', 'three']:
                await pub.publish('channel:1', msg)
            # send stop word
            await pub.publish('channel:1', STOPWORD)
        pub.close()
        await pub.wait_closed()

    loop.run_until_complete(asyncio.gather(publish(), tsk, loop=loop)) 

def __getattr__(self, name):
        assert not self._done, "Pipeline already executed. Create new one."
        attr = getattr(self._redis, name)
        if callable(attr):

            @functools.wraps(attr)
            def wrapper(*args, **kw):
                try:
                    task = asyncio.ensure_future(attr(*args, **kw))
                except Exception as exc:
                    task = get_event_loop().create_future()
                    task.set_exception(exc)
                self._results.append(task)
                return task
            return wrapper
        return attr 

def _fail(self, e: Exception, clean_close: bool = True) -> None:
        if self.close_task is None:
            self.close_task = asyncio.shield(
                asyncio.ensure_future(self._close_coro(e, clean_close=clean_close))
            ) 

def test_websocket_subscription_task_cancel(
    event_loop, client_and_server, subscription_str
):

    session, server = client_and_server

    count = 10
    subscription = gql(subscription_str.format(count=count))

    async def task_coro():
        nonlocal count
        async for result in session.subscribe(subscription):

            number = result["number"]
            print(f"Number received: {number}")

            assert number == count

            count -= 1

    task = asyncio.ensure_future(task_coro())

    async def cancel_task_coro():
        nonlocal task

        await asyncio.sleep(11 * MS)

        task.cancel()

    cancel_task = asyncio.ensure_future(cancel_task_coro())

    await asyncio.gather(task, cancel_task)

    assert count > 0 

def test_websocket_subscription_close_transport(
    event_loop, client_and_server, subscription_str
):

    session, server = client_and_server

    count = 10
    subscription = gql(subscription_str.format(count=count))

    async def task_coro():
        nonlocal count
        async for result in session.subscribe(subscription):

            number = result["number"]
            print(f"Number received: {number}")

            assert number == count

            count -= 1

    task = asyncio.ensure_future(task_coro())

    async def close_transport_task_coro():
        nonlocal task

        await asyncio.sleep(11 * MS)

        await session.transport.close()

    close_transport_task = asyncio.ensure_future(close_transport_task_coro())

    await asyncio.gather(task, close_transport_task)

    assert count > 0 

def test_websocket_two_queries_in_parallel(
    event_loop, client_and_server, query_str
):

    session, server = client_and_server

    query = gql(query_str)

    result1 = None
    result2 = None

    async def task1_coro():
        nonlocal result1
        result1 = await session.execute(query)

    async def task2_coro():
        nonlocal result2
        result2 = await session.execute(query)

    task1 = asyncio.ensure_future(task1_coro())
    task2 = asyncio.ensure_future(task2_coro())

    await asyncio.gather(task1, task2)

    print("Query1 received:", result1)
    print("Query2 received:", result2)

    assert result1 == result2 

def main():
  future = asyncio.Future()
  await asyncio.ensure_future(myFuture(future))
  print(future.result()) 

def myGenerator():
   for i in range(5):
       asyncio.ensure_future(myTask(i))
   print("Completed Tasks")
   yield from asyncio.sleep(2) 

def send_code(self, request):
        uid = int(await request.text())
        if uid in self._uid_to_code.keys():
            return web.Response()
        code = secrets.randbelow(100000)
        asyncio.ensure_future(asyncio.shield(self._clear_code(uid)))
        self._uid_to_code[uid] = b64encode(hashlib.scrypt((str(code).zfill(5) + str(uid)).encode("utf-8"),
                                                          salt="friendlytgbot".encode("utf-8"),
                                                          n=16384, r=8, p=1, dklen=64)).decode("utf-8")
        await self.client_data[uid][1].send_message("me", "Your code is <code>{:05d}</code>\nDo <b>not</b> "
                                                          "share this code with anyone, even is they say they are"
                                                          " from friendly-telegram.\nThe code will expire in "
                                                          "2 minutes.".format(code))
        return web.Response() 

def check_code(self, request):
        code, uid = (await request.text()).split("\n")
        uid = int(uid)
        if uid not in self._uid_to_code:
            return web.Response(status=404)
        if self._uid_to_code[uid] == code:
            del self._uid_to_code[uid]
            secret = secrets.token_urlsafe()
            asyncio.ensure_future(asyncio.shield(self._clear_secret(secret)))
            self._secret_to_uid[secret] = uid  # If they just signed in, they automatically are authenticated
            return web.Response(text=secret)
        else:
            return web.Response(status=401) 

def save(self):
        if self._pending is not None and not self._pending.cancelled():
            self._pending.cancel()
        if self._sync_future is None or self._sync_future.done():
            self._sync_future = NotifyingFuture(on_await=self._cancel_then_set)
        self._pending = asyncio.ensure_future(_wait_then_do(10, self._set))  # Delay database ops by 10s
        return self._sync_future 

def _cancel_then_set(self):
        if self._pending is not None and not self._pending.cancelled():
            self._pending.cancel()
        self._pending = asyncio.ensure_future(self._set())
        # Restart the task, but without the delay, because someone is waiting for us 

def listen_message_stream(self, id_blacklist=None):
        id_blacklist = set(id_blacklist or [self.me, ])

        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        with aiohttp.ClientSession(loop=loop) as session:
            self.aioclient_session = session

            tasks = [
                asyncio.ensure_future(self.fetch(session, room, id_blacklist))
                for room in self.rooms
            ]
            done, _ = loop.run_until_complete(
                asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION)
            )
            for d in done:
                if d.exception():
                    raise d.exception()