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 __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 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 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 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 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 fetch_bar():
    day = datetime.datetime.strptime('20100416', '%Y%m%d').replace(tzinfo=pytz.FixedOffset(480))
    end = datetime.datetime.strptime('20160118', '%Y%m%d').replace(tzinfo=pytz.FixedOffset(480))
    tasks = []
    while day <= end:
        tasks.append(is_trading_day(day))
        day += datetime.timedelta(days=1)
    trading_days = []
    for f in tqdm(asyncio.as_completed(tasks), total=len(tasks)):
        rst = await f
        trading_days.append(rst)
    tasks.clear()
    for day, trading in trading_days:
        if trading:
            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)),
            ]
    print('task len=', len(tasks))
    for f in tqdm(asyncio.as_completed(tasks), total=len(tasks)):
        await f 

def query(self, query_type: str, **kwargs):
        sub_client = None
        channel_name1, channel_name2 = None, None
        try:
            sub_client = await aioredis.create_redis(
                (config.get('REDIS', 'host', fallback='localhost'),
                 config.getint('REDIS', 'port', fallback=6379)),
                db=config.getint('REDIS', 'db', fallback=1))
            request_id = self.next_id()
            kwargs['RequestID'] = request_id
            channel_name1 = self.__trade_response_format.format('OnRspQry' + query_type, request_id)
            channel_name2 = self.__trade_response_format.format('OnRspError', request_id)
            ch1, ch2 = await sub_client.psubscribe(channel_name1, channel_name2)
            cb = self.io_loop.create_future()
            tasks = [
                asyncio.ensure_future(self.query_reader(ch1, cb), loop=self.io_loop),
                asyncio.ensure_future(self.query_reader(ch2, cb), loop=self.io_loop),
            ]
            self.redis_client.publish(self.__request_format.format('ReqQry' + query_type), json.dumps(kwargs))
            rst = await asyncio.wait_for(cb, HANDLER_TIME_OUT, loop=self.io_loop)
            await sub_client.punsubscribe(channel_name1, channel_name2)
            sub_client.close()
            await asyncio.wait(tasks, loop=self.io_loop)
            return rst
        except Exception as e:
            logger.error('%s failed: %s', query_type, repr(e), exc_info=True)
            if sub_client and sub_client.in_pubsub and channel_name1:
                await sub_client.unsubscribe(channel_name1, channel_name2)
                sub_client.close()
            return None 

def get_data(self, raw: bool = True) -> AnyStr:
        """The request body data."""
        try:
            body_future = asyncio.ensure_future(self.body)
            raw_data = await asyncio.wait_for(body_future, timeout=self.body_timeout)
        except asyncio.TimeoutError:
            body_future.cancel()
            from ..exceptions import RequestTimeout  # noqa Avoiding circular import

            raise RequestTimeout()

        if raw:
            return raw_data
        else:
            return raw_data.decode(self.charset) 

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 main():
    sock = socket.socket()
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind((ip, port))
    sock.listen(1)
    logging.debug('Listening')
    conn, _ = sock.accept()
    logging.debug('Connected')

    while True:
        data = json.loads(conn.recv(chunk_size).decode("utf-8"))
        logging.debug('received data: {}'.format(data))
        if ('type' in data) & (data['type'] == 'get_claim_def'):
            logging.debug('get_claim_def -> start')
            init = asyncio.ensure_future(issuer_init(primes1(), conn))
            loop.run_until_complete(init)
            logging.debug('get_claim_def -> done')
        if ('type' in data) & (data['type'] == 'issue_claim'):
            logging.debug('issue_claim -> start')
            future = asyncio.ensure_future(
                issue_claim(conn, data['data']['blinded_ms']))
            loop.run_until_complete(future)
            logging.debug('issue_claim -> done')
        if (('type' in data) & (data['type'] == 'close')) | (not data):
            break

    sock.close() 

def _list_sdcard(self, done_cb):
        asyncio.ensure_future(self._parse_sdcard_list(done_cb)) 

def _stream_file(self, fn):
        self.file_streamer = asyncio.ensure_future(self.stream_file(fn)) 

def _upload_gcode(self, fn, donecb):
        self.file_streamer = asyncio.ensure_future(self._stream_upload_gcode(fn, donecb)) 

def __init__(self, cb):
        super().__init__()
        self.cb = cb
        self.cnt= 0
        self.log = logging.getLogger() #.getChild('TcpConnection')
        self.log.info('TcpConnection: creating TcpCOnnection')
        self.queue = asyncio.Queue(maxsize=100)
        self.hipri_queue = asyncio.Queue()
        self._ready = asyncio.Event()
        self._msg_ready = asyncio.Semaphore(value=0)
        self.tsk= asyncio.async(self._send_messages())  # Or asyncio.ensure_future if using 3.4.3+
        self.flush= False 

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

        expected_response = mocker.Mock(spec=aiohttp_.aiohttp.ClientResponse)

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

        aiohttp_session_mock.request = request
        client = aiohttp_.AiohttpClient(aiohttp_session_mock)
        client._sync_callback_adapter = asyncio.coroutine

        # Run
        globals_ = dict(globals(), client=client)
        locals_ = {"asyncio": asyncio}
        exec(
            """@asyncio.coroutine
def call():
    response = yield from client.send((1, 2, {}))
    response = yield from client.apply_callback(lambda x: 2, response)
    return response
""",
            globals_,
            locals_,
        )

        loop = asyncio.get_event_loop()
        value = loop.run_until_complete(
            asyncio.ensure_future(locals_["call"]())
        )

        # Verify
        assert value == 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()