import asyncio


async def main():
    loop = asyncio.get_event_loop()
    f = asyncio.Future()

    def _on_complete():
        f.set_result(True)

    loop.call_later(1, _on_complete)
    return await f


if __name__ == "__main__":
    import time
    s_t = time.time()
    asyncio.run(main())
    print(time.time() - s_t)

@types.coroutine
def __sleep0():
    yield


async def sleep(delay, result=None):
    """Coroutine that completes after a given time (in seconds)."""
    if delay <= 0:
        await __sleep0()
        return result

    loop = events.get_running_loop()
    future = loop.create_future()
    h = loop.call_later(delay,
                        futures._set_result_unless_cancelled,
                        future, result)
    try:
        return await future
    finally:
        h.cancel()

import asyncio

async def demo() -> None:
    for index, i in enumerate(range(10000)):
        if index % 100 == 0:
            await asyncio.sleep(0)
        ...  # 假設(shè)這里的代碼占用過多的CPU時間

import asyncio


async def sub(f):
    await asyncio.shield(f)


async def main():
    f1 = asyncio.Future()
    f2 = asyncio.Future()
    sub1 = asyncio.create_task(sub(f1))
    sub2 = asyncio.create_task(sub(f2))
    f1.cancel()
    sub2.cancel()
    await asyncio.sleep(0)  # 確保已經(jīng)取消完成
    print("f1 future run success:", f1.done())
    print("f2 future run success:", f2.done())
    print("sub1 future run result:", sub1.done())
    print("sub2 future run result:", sub2.done())

asyncio.run(main())

# >>> future run success: True
# >>> future run success: False
# >>> sub1 future run result: True
# >>> sub2 future run result: True

def shield(arg):
    # 如果是Coro，則需要包裝成future
    inner = _ensure_future(arg)
    if inner.done():
        # 如果已經(jīng)完成，就不需要被處理了
        return inner
    loop = futures._get_loop(inner)
    # 創(chuàng)建一個future容器
    outer = loop.create_future()

    def _inner_done_callback(inner):
        if outer.cancelled():
            if not inner.cancelled():
                # 如果容器已經(jīng)被取消，而自己沒被取消且已經(jīng)完成，則手動獲取下結(jié)果，方便被回收
                inner.exception()
            return

        if inner.cancelled():
            # 如果自己被取消，則把取消通過容器傳播到協(xié)程鏈上
            outer.cancel()
        else:
            # 自己已經(jīng)完成且容器未完成，把自己的結(jié)果或者異常通過替身傳播到協(xié)程鏈上
            exc = inner.exception()
            if exc is not None:
                outer.set_exception(exc)
            else:
                outer.set_result(inner.result())


    def _outer_done_callback(outer):
        if not inner.done():
            inner.remove_done_callback(_inner_done_callback)

    # 添加回調(diào)，在執(zhí)行成功或被取消時通知對方
    inner.add_done_callback(_inner_done_callback)
    outer.add_done_callback(_outer_done_callback)
    return outer

def _release_waiter(waiter, *args):
    if not waiter.done():
        waiter.set_result(None)

async def _cancel_and_wait(fut, loop):
    waiter = loop.create_future()
    cb = functools.partial(_release_waiter, waiter)
    fut.add_done_callback(cb)

    try:
        fut.cancel()
        await waiter
    finally:
        fut.remove_done_callback(cb)

async def wait_for(fut, timeout):
    loop = events.get_running_loop()

    if timeout is None:
        return await fut

    if timeout <= 0:
        # 當(dāng)超時的值小于等于0時就意味著想馬上得到結(jié)果
        
        fut = ensure_future(fut, loop=loop)

        if fut.done():
            # 如果執(zhí)行完成就返回可等待對象的數(shù)據(jù)
            return fut.result()
        # 取消可等待對象并等待
        await _cancel_and_wait(fut, loop=loop)
        # 如果被_cancel_and_wait取消，那么會拋出CancelledError異常，這時候把它轉(zhuǎn)為超時異常
        try:
            return fut.result()
        except exceptions.CancelledError as exc:
            raise exceptions.TimeoutError() from exc

    # 初始化一個Future,只有在超時和完成時才會變?yōu)閐one
    waiter = loop.create_future()
    timeout_handle = loop.call_later(timeout, _release_waiter, waiter)
    cb = functools.partial(_release_waiter, waiter)

    fut = ensure_future(fut, loop=loop)
    fut.add_done_callback(cb)

    try:
        try:
            await waiter
        except exceptions.CancelledError:
            # 此時是asyncio.Task被取消，并把取消傳播到waiter
            if fut.done():
                return fut.result()
            else:
                # 如果任務(wù)被取消了，那么需要確保任務(wù)沒有被執(zhí)行才返回
                fut.remove_done_callback(cb)
                await _cancel_and_wait(fut, loop=loop)
                raise
        # 計時結(jié)束或者是執(zhí)行完畢的情況
        if fut.done():
            # 執(zhí)行完畢，返回對應(yīng)的值
            return fut.result()
        else:
            # 計時結(jié)束，清理資源，并拋出異常
            fut.remove_done_callback(cb)
            # 如果任務(wù)被取消了，那么需要確保任務(wù)沒有被執(zhí)行才返回
            await _cancel_and_wait(fut, loop=loop)
            # 如果被_cancel_and_wait取消，那么會拋出CancelledError異常，這時候把它轉(zhuǎn)為超時異常
            try:
                return fut.result()
            except exceptions.CancelledError as exc:
                raise exceptions.TimeoutError() from exc
    finally:
        timeout_handle.cancel()

import asyncio


async def main():
    return await asyncio.wait(
        {asyncio.create_task(asyncio.sleep(1))},
        timeout=0.5
    )


if __name__ == "__main__":
    asyncio.run(main())

async def _wait(fs, timeout, return_when, loop):
    assert fs, 'Set of Futures is empty.'
    waiter = loop.create_future()
    timeout_handle = None
    if timeout is not None:
        # 定義一個time handler,在timeout秒后通過`_release_waiter`完成.
        timeout_handle = loop.call_later(timeout, _release_waiter, waiter)
    counter = len(fs)

    def _on_completion(f):
        # 每個可等待對象執(zhí)行完成的回調(diào)
        nonlocal counter
        counter -= 1
        if (counter <= 0 or
            return_when == FIRST_COMPLETED or
            return_when == FIRST_EXCEPTION and
             (not f.cancelled() and f.exception() is not None)
        ):
            # 如果所有任務(wù)執(zhí)行完成，或者是第一個完成或者是第一個拋出異常時，
            # 意味著執(zhí)行完成，需要取消time handler,并標(biāo)記為完成
            if timeout_handle is not None:
                timeout_handle.cancel()
            if not waiter.done():
                waiter.set_result(None)
    # 為每個可等待對象添加回調(diào)
    for f in fs:
        f.add_done_callback(_on_completion)

    try:
        # 等待替身執(zhí)行完成
        await waiter
    finally:
        # 取消time handler并移除回調(diào)(因為cancel是異步的)
        if timeout_handle is not None:
            timeout_handle.cancel()
        for f in fs:
            f.remove_done_callback(_on_completion)

    # 處理并返回done和pending，其中done代表完成，pending代表執(zhí)行中。
    done, pending = set(), set()
    for f in fs:
        if f.done():
            done.add(f)
        else:
            pending.add(f)
    return done, pending

import asyncio


async def main():
    f_list = [asyncio.Future() for _ in range(10)]
    done, pending = await asyncio.wait(f_list, timeout=1)
    print(len(done), len(pending))
    print([i for i in pending if i.done()])
    f_list[1].set_result(True)
    print([i for i in pending if i.done()])


if __name__ == "__main__":
    asyncio.run(main())
# >>> 0 10
# >>> []
# >>> [<Future finished result=True>]

import asyncio


async def sub(i):
    await asyncio.sleep(i)
    return i


async def main():
    for f in asyncio.as_completed([sub(i) for i in range(5)], timeout=3):
        print(await f)


if __name__ == "__main__":
    asyncio.run(main())
# >>> 0
# >>> 1
# >>> 2
# >>> Traceback (most recent call last):
#       File "/home/so1n/github/demo_project/demo.py", line 18, in <module>
#         asyncio.run(main())
#       File "/usr/lib/python3.7/asyncio/runners.py", line 43, in run
#         return loop.run_until_complete(main)
#       File "/usr/lib/python3.7/asyncio/base_events.py", line 584, in run_until_complete
#         return future.result()
#       File "/home/so1n/github/demo_project/demo.py", line 14, in main
#         print(await f)
#       File "/usr/lib/python3.7/asyncio/tasks.py", line 532, in _wait_for_one
#         raise futures.TimeoutError
#     concurrent.futures._base.TimeoutError

def as_completed(fs, *, timeout=None):
    from .queues import Queue  # Import here to avoid circular import problem.
    done = Queue()

    loop = events._get_event_loop()
    todo = {ensure_future(f, loop=loop) for f in set(fs)}
    timeout_handle = None

    def _on_timeout():
        # 超時時調(diào)用，需要注意的是，失敗時結(jié)果為空，所以要推送一個空的數(shù)據(jù)到隊列中
        # 在消費者發(fā)現(xiàn)元素為空時拋出錯誤
        for f in todo:
            f.remove_done_callback(_on_completion)
            done.put_nowait(None)  # Queue a dummy value for _wait_for_one().
        todo.clear()  # Can't do todo.remove(f) in the loop.

    def _on_completion(f):
        # 如果成功，就把Future推送到隊列中，消費者可以通過Future獲取到結(jié)果
        if not todo:
            return  # _on_timeout() was here first.
        todo.remove(f)
        done.put_nowait(f)
        if not todo and timeout_handle is not None:
            timeout_handle.cancel()

    async def _wait_for_one():
        f = await done.get()
        if f is None:
            # 如果元素為空，則證明已經(jīng)超時了，要拋出異常
            raise exceptions.TimeoutError
        return f.result()

    for f in todo:
        f.add_done_callback(_on_completion)
    if todo and timeout is not None:
        timeout_handle = loop.call_later(timeout, _on_timeout)
    # 通過生成器語法返回協(xié)程函數(shù)，該協(xié)程函數(shù)可以獲取最近完成的可等待對象的結(jié)果
    for _ in range(len(todo)):
        yield _wait_for_one()

import asyncio
import random


async def sub():
    # 一半的幾率會被set一個值并返回，一半的幾率會卡死
    f = asyncio.Future()
    if random.choice([0, 1]) == 0:
        f.set_result(None)
    return await f


async def main():
    try:
        for f in asyncio.as_completed([sub() for i in range(5)], timeout=1):
            print(await f)
    except asyncio.TimeoutError:
        # 忽略超時
        pass
    # 統(tǒng)計未完成的sub任務(wù)
    cnt = 0
    for i in asyncio.all_tasks():
        if i._coro.__name__ == sub.__name__:
            cnt += 1
    print("runing task by name sub:", cnt)


if __name__ == "__main__":
    asyncio.run(main())
# >>> None
# >>> None
# >>> None
# >>> runing task by name sub: 2