转-python异步IO-asyncio

import argparse
import asyncio
import functools
import logging
import random
import urllib.parse

loop = asyncio.get_event_loop()

@asyncio.coroutine
def print_http_headers(no, url, keepalive):
    url = urllib.parse.urlsplit(url)
    wait_for = functools.partial(asyncio.wait_for, timeout=3, loop=loop)
    query = (‘HEAD {url.path} HTTP/1.1\r\n‘
             ‘Host: {url.hostname}\r\n‘
             ‘\r\n‘).format(url=url).encode(‘utf-8‘)

    rd, wr = yield from wait_for(asyncio.open_connection(url.hostname, 80))
    while True:
        wr.write(query)
   
        while True:
            line = yield from wait_for(rd.readline())
            if not line: # end of connection
                wr.close()
                return no
            line = line.decode(‘utf-8‘).rstrip()
            if not line: # end of header
                break
            logging.debug(‘(%d) HTTP header> %s‘ % (no, line))

        yield from asyncio.sleep(random.randint(1, keepalive//2))


@asyncio.coroutine
def do_requests(args):
    conn_pool = set()
    waiter = asyncio.Future()

    def _on_complete(fut):
        conn_pool.remove(fut)
        exc, res = fut.exception(), fut.result()
        if exc is not None:
            logging.info(‘conn#{} exception‘.format(exc))
        else:
            logging.info(‘conn#{} result‘.format(res))

        if not conn_pool:
            waiter.set_result(‘event loop is done‘)

    for i in range(args.connections):
        fut = asyncio.async(print_http_headers(i, args.url, args.keepalive))
        fut.add_done_callback(_on_complete)
        conn_pool.add(fut)
        if i % 10 == 0:
            yield from asyncio.sleep(0.01)

    logging.info((yield from waiter))


def main():
    parser = argparse.ArgumentParser(description=‘asyncli‘)
    parser.add_argument(‘url‘, help=‘page address‘)
    parser.add_argument(‘-c‘, ‘--connections‘, type=int, default=1,
                        help=‘number of connections simultaneously‘)
    parser.add_argument(‘-k‘, ‘--keepalive‘, type=int, default=60,
                        help=‘HTTP keepalive timeout‘)
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format=‘%(asctime)s %(message)s‘)
    loop.run_until_complete(do_requests(args))
    loop.close()


if __name__ == ‘__main__‘:
    main()