asyncio系列之Queue实现

import types
import select
import time
import socket
import functools
import collections


class Future:
    def __init__(self, *, loop=None):
        self._result = None
        self._callbacks = []
        self._loop = loop

    def set_result(self, result):
        self._result = result
        callbacks = self._callbacks[:]
        self._callbacks = []
        for callback in callbacks:
            self._loop._ready.append(callback)

    def add_callback(self, callback):
        self._callbacks.append(callback)

    def __iter__(self):
        print("挂起在yield处")
        yield self
        print("恢复执行")
        return "future"

    __await__ = __iter__


class Task:
    def __init__(self, cor, *, loop=None):
        self.cor = cor
        self._loop = loop

    def _step(self):
        cor = self.cor
        try:
            result = cor.send(None)
        except StopIteration as e:
            self._loop._task_count -= 1
            if self._loop._task_count == 0:
                self._loop.close()
        except Exception as e:
            pass
        else:
            if isinstance(result, Future):
                result.add_callback(self._wakeup)

    def _wakeup(self):
        self._step()


class Loop:
    def __init__(self):
        self._stop = False
        self._ready = []
        self._scheduled = []
        self._time = lambda: time.time()
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.setblocking(False)
        self._select = functools.partial(select.select, [sock], [], [])
        self._task_count = 0

    def create_task(self, cor):
        task = Task(cor, loop=self)
        self._ready.append(task._step)
        self._task_count += 1
        return task

    def call_later(self, delay, callback, *args):
        callback._when = delay
        self._scheduled.append((callback, *args))

    def run_until_complete(self, task):
        assert isinstance(task, Task)
        timeout = None
        while not self._stop:
            if self._ready:
                timeout = 0
            if self._scheduled:
                callback, *args = self._scheduled.pop()
                timeout = callback._when
                self._ready.append(functools.partial(callback, *args))

                self._select(timeout)
            n = len(self._ready)
            for i in range(n):
                step = self._ready.pop()
                step()

    def close(self):
        self._stop = True


@types.coroutine
def _sleep():
    yield


async def sleep(s, result=None):
    if s <= 0:
        await _sleep()
        return result
    else:
        future = Future(loop=loop)
        future._loop.call_later(s, unless_cancelled, future)
        await future
        return result


def unless_cancelled(future):
    future.set_result(None)


class Queue:
    def __init__(self, maxsize=0, *, loop=None):
        self._loop = loop
        self._maxsize = maxsize

        # Futures.
        self._getters = collections.deque()
        # Futures.
        self._putters = collections.deque()
        self._init(maxsize)

    def _init(self, maxsize):
        self._queue = collections.deque()

    def _get(self):
        return self._queue.popleft()

    def _put(self, item):
        self._queue.append(item)

    def _wakeup_next(self, waiters):
        while waiters:
            waiter = waiters.popleft()
            try:
                waiter.set_result(None)
            except Exception as e:
                pass
            break

    def qsize(self):
        return len(self._queue)

    @property
    def maxsize(self):
        return self._maxsize

    def empty(self):
        return not self._queue

    def full(self):
        if self._maxsize <= 0:
            return False
        else:
            return self.qsize() >= self._maxsize

    @types.coroutine
    def put(self, item):
        while self.full():
            putter = Future(loop=self._loop)
            self._putters.append(putter)
            try:
                yield from putter
            except:
                if not self.full():
                    self._wakeup_next(self._putters)
                raise
        return self.put_nowait(item)

    def put_nowait(self, item):
        if self.full():
            raise
        self._put(item)
        self._wakeup_next(self._getters)

    @types.coroutine
    def get(self):
        while self.empty():
            getter = Future(loop=self._loop)
            self._getters.append(getter)
            try:
                yield from getter
            except:
                try:
                    self._getters.remove(getter)
                except ValueError:
                    pass

                if not self.empty():
                    self._wakeup_next(self._getters)
                raise
        return self.get_nowait()

    def get_nowait(self):
        if self.empty():
            raise
        item = self._get()
        self._wakeup_next(self._putters)
        return item


async def foo(queue):
    print(f'enter foo at {time.strftime("%Y-%m-%d %H:%M:%S")}')
    item = await queue.get()
    print('foo get ', item)
    await sleep(2)
    item = await queue.get()
    print('foo get ', item)
    print(f'exit foo at {time.strftime("%Y-%m-%d %H:%M:%S")}')


async def goo(queue):
    print(f'enter goo at {time.strftime("%Y-%m-%d %H:%M:%S")}')
    print('goo put a')
    await queue.put('a')
    # await sleep(2)
    print('goo put b')
    await queue.put('b')
    print(f'exit goo at {time.strftime("%Y-%m-%d %H:%M:%S")}')

if __name__ == '__main__':
    loop = Loop()
    queue = Queue(maxsize=1, loop=loop)
    f = foo(queue)
    g = goo(queue)
    task1 = loop.create_task(f)
    task2 = loop.create_task(g)
    loop.run_until_complete(task1)