PyWebIO/pywebio/session/coroutinebased.py

import asyncio
import logging
import threading
from contextlib import contextmanager
from functools import partial

from .base import Session
from ..exceptions import SessionNotFoundException, SessionClosedException, SessionException
from ..utils import random_str, isgeneratorfunction, iscoroutinefunction

logger = logging.getLogger(__name__)


class WebIOFuture:
    def __init__(self, coro=None):
        self.coro = coro

    def __iter__(self):
        result = yield self
        return result

    __await__ = __iter__  # make compatible with 'await' expression


class _context:
    current_session = None  # type:"CoroutineBasedSession"
    current_task_id = None


class CoroutineBasedSession(Session):
    """
    基于协程的任务会话

    当主协程任务和会话内所有通过 `run_async` 注册的协程都退出后，会话关闭。
    当用户浏览器主动关闭会话，CoroutineBasedSession.close 被调用， 协程任务和会话内所有通过 `run_async` 注册的协程都被关闭。

    """

    # 运行事件循环的线程id
    # 用于在 CoroutineBasedSession.get_current_session() 判断调用方是否合法
    # Tornado backend时，在创建第一个CoroutineBasedSession时初始化
    # Flask backend时，在platform.flaskrun_event_loop()时初始化
    event_loop_thread_id = None

    @classmethod
    def get_current_session(cls) -> "CoroutineBasedSession":
        if _context.current_session is None or cls.event_loop_thread_id != threading.current_thread().ident:
            raise SessionNotFoundException("No session found in current context!")

        if _context.current_session.closed():
            raise SessionClosedException

        return _context.current_session

    @staticmethod
    def get_current_task_id():
        if _context.current_task_id is None:
            raise RuntimeError("No current task found in context!")
        return _context.current_task_id

    def __init__(self, target, session_info, on_task_command=None, on_session_close=None):
        """
        :param target: 协程函数
        :param on_task_command: 由协程内发给session的消息的处理函数
        :param on_session_close: 会话结束的处理函数。后端Backend在相应on_session_close时关闭连接时，需要保证会话内的所有消息都传送到了客户端
        """
        assert iscoroutinefunction(target) or isgeneratorfunction(target), ValueError(
            "CoroutineBasedSession accept coroutine function or generator function as task function")

        super().__init__(session_info)

        cls = type(self)

        self._on_task_command = on_task_command or (lambda _: None)
        self._on_session_close = on_session_close or (lambda: None)

        # 当前会话未被Backend处理的消息
        self.unhandled_task_msgs = []

        # 在创建第一个CoroutineBasedSession时 event_loop_thread_id 还未被初始化
        # 则当前线程即为运行 event loop 的线程
        if cls.event_loop_thread_id is None:
            cls.event_loop_thread_id = threading.current_thread().ident

        # 会话内的协程任务
        self.coros = {}  # coro_task_id -> Task()

        self._closed = False

        self._need_keep_alive = False

        # 当前会话未结束运行(已创建和正在运行的)的协程数量。当 _alive_coro_cnt 变为 0 时，会话结束。
        self._alive_coro_cnt = 1

        main_task = Task(self._start_main_task(target), session=self, on_coro_stop=self._on_task_finish)
        self.coros[main_task.coro_id] = main_task

        self._step_task(main_task)

    async def _start_main_task(self, target):
        await target()
        if self.need_keep_alive():
            from ..session import hold
            await hold()

    def _step_task(self, task, result=None):
        asyncio.get_event_loop().call_soon_threadsafe(partial(task.step, result))

    def _on_task_finish(self, task: "Task"):
        self._alive_coro_cnt -= 1

        if task.coro_id in self.coros:
            logger.debug('del self.coros[%s]', task.coro_id)
            del self.coros[task.coro_id]

        if self._alive_coro_cnt <= 0 and not self.closed():
            self.send_task_command(dict(command='close_session'))
            self._on_session_close()
            self.close()

    def send_task_command(self, command):
        """向会话发送来自协程内的消息

        :param dict command: 消息
        """
        if self.closed():
            raise SessionClosedException()
        self.unhandled_task_msgs.append(command)
        self._on_task_command(self)

    async def next_client_event(self):
        # 函数开始不需要判断 self.closed()
        # 如果会话关闭，对 get_current_session().next_client_event() 的调用会抛出SessionClosedException
        return await WebIOFuture()

    def send_client_event(self, event):
        """向会话发送来自用户浏览器的事件️

        :param dict event: 事件️消息
        """
        coro_id = event['task_id']
        coro = self.coros.get(coro_id)
        if not coro:
            logger.error('coro not found, coro_id:%s', coro_id)
            return
        self._step_task(coro, event)

    def get_task_commands(self):
        msgs = self.unhandled_task_msgs
        self.unhandled_task_msgs = []
        return msgs

    def _cleanup(self):
        for t in list(self.coros.values()):  # t.close() may cause self.coros changed size
            t.step(SessionClosedException, throw_exp=True)
            # in case that the task catch the SessionClosedException, we need to close it manually
            t.close()
        self.coros = {}  # delete session tasks

        # reset the reference, to avoid circular reference
        self._on_session_close = None
        self._on_task_command = None

    def close(self, nonblock=False):
        """关闭当前Session。由Backend调用"""
        if self.closed():
            return

        super().close()

        self._cleanup()

    def register_callback(self, callback, mutex_mode=False):
        """ 向Session注册一个回调函数，返回回调id

        :type callback: Callable or Coroutine
        :param callback: 回调函数. 函数签名为 ``callback(data)``. ``data`` 参数为回调事件的值
        :param bool mutex_mode: 互斥模式。若为 ``True`` ，则在运行回调函数过程中，无法响应同一组件（callback_id相同）的新点击事件，仅当 ``callback`` 为协程函数时有效
        :return str: 回调id.
        """

        async def callback_coro():
            while True:
                try:
                    event = await self.next_client_event()
                except SessionClosedException:
                    return

                assert event['event'] == 'callback'
                coro = None
                if iscoroutinefunction(callback):
                    coro = callback(event['data'])
                elif isgeneratorfunction(callback):
                    coro = asyncio.coroutine(callback)(event['data'])
                else:
                    try:
                        res = callback(event['data'])
                        if asyncio.iscoroutine(res):
                            coro = res
                        else:
                            del res  # `res` maybe pywebio.io_ctrl.Output, so need release `res`
                    except Exception:
                        self.on_task_exception()

                if coro is not None:
                    if mutex_mode:
                        await coro
                    else:
                        self.run_async(coro)

        cls = type(self)
        callback_task = Task(callback_coro(), cls.get_current_session())
        # Activate task
        # Don't callback.step(), it will result in recursive calls to step()
        # todo: integrate with inactive_coro_instances
        callback_task.coro.send(None)
        cls.get_current_session().coros[callback_task.coro_id] = callback_task

        self._need_keep_alive = True

        return callback_task.coro_id

    def run_async(self, coro_obj):
        """异步运行协程对象。可以在协程内调用 PyWebIO 交互函数

        :param coro_obj: 协程对象
        :return: An instance of  `TaskHandler` is returned, which can be used later to close the task.
        """
        assert asyncio.iscoroutine(coro_obj), '`run_async()` only accept coroutine object'

        self._alive_coro_cnt += 1

        task = Task(coro_obj, session=self, on_coro_stop=self._on_task_finish)
        self.coros[task.coro_id] = task
        asyncio.get_event_loop().call_soon_threadsafe(task.step)
        return task.task_handle()

    async def run_asyncio_coroutine(self, coro_obj):
        """若会话线程和运行事件的线程不是同一个线程，需要用 asyncio_coroutine 来运行asyncio中的协程"""
        assert asyncio.iscoroutine(coro_obj), '`run_asyncio_coroutine()` only accept coroutine object'

        res = await WebIOFuture(coro=coro_obj)
        return res

    def need_keep_alive(self) -> bool:
        return self._need_keep_alive


class TaskHandler:
    """The handler of coroutine task

    See also: `run_async() <pywebio.session.run_async>`
    """

    def __init__(self, close, closed):
        self._close = close
        self._closed = closed

    def close(self):
        """Close the coroutine task."""
        return self._close()

    def closed(self) -> bool:
        """Returns a bool stating whether the coroutine task is closed. """
        return self._closed()


class Task:

    @contextmanager
    def session_context(self):
        """
        >>> with session_context():
        ...     res = self.coros[-1].send(data)
        """

        # todo issue: with 语句可能发生嵌套，导致内层with退出时，将属性置空
        _context.current_session = self.session
        _context.current_task_id = self.coro_id
        try:
            yield
        finally:
            _context.current_session = None
            _context.current_task_id = None

    @staticmethod
    def gen_coro_id(coro=None):
        """生成协程id"""
        name = 'coro'
        if hasattr(coro, '__name__'):
            name = coro.__name__

        return '%s-%s' % (name, random_str(10))

    def __init__(self, coro, session: CoroutineBasedSession, on_coro_stop=None):
        """
        :param coro: 协程对象
        :param session: 创建该Task的会话实例
        :param on_coro_stop: 任务结束(正常结束或外部调用Task.close)时运行的回调
        """
        self.session = session
        self.coro = coro
        self.result = None
        self.task_closed = False  # 任务完毕/取消
        self.on_coro_stop = on_coro_stop or (lambda _: None)

        self.coro_id = self.gen_coro_id(self.coro)

        self.pending_futures = {}  # id(future) -> future

        logger.debug('Task[%s] created ', self.coro_id)

    def step(self, result=None, throw_exp=False):
        """激活协程

        :param any result: 向协程传入的数据
        :param bool throw_exp: 是否向协程引发异常，为 True 时， result 参数为相应的异常对象
        """
        coro_yield = None
        with self.session_context():
            try:
                if throw_exp:
                    coro_yield = self.coro.throw(result)
                else:
                    coro_yield = self.coro.send(result)
            except StopIteration as e:
                if len(e.args) == 1:
                    self.result = e.args[0]
                self.close()
                logger.debug('Task[%s] finished', self.coro_id)
            except Exception as e:
                if not isinstance(e, SessionException):
                    self.session.on_task_exception()
                self.close()

        if coro_yield is None:
            return

        future = None
        if isinstance(coro_yield, WebIOFuture):
            if coro_yield.coro:
                future = asyncio.run_coroutine_threadsafe(coro_yield.coro, asyncio.get_event_loop())
        else:
            future = coro_yield

        if not self.session.closed() and hasattr(future, 'add_done_callback'):
            future.add_done_callback(self._wakeup)
            self.pending_futures[id(future)] = future

    def _wakeup(self, future):
        if not future.cancelled():
            del self.pending_futures[id(future)]
            self.step(future.result())

    def close(self):
        if self.task_closed:
            return

        self.task_closed = True

        self.coro.close()
        while self.pending_futures:
            _, f = self.pending_futures.popitem()
            f.cancel()

        self.on_coro_stop(self)
        self.on_coro_stop = None  # avoid circular reference
        self.session = None

        logger.debug('Task[%s] closed', self.coro_id)

    def __del__(self):
        if not self.task_closed:
            logger.warning('Task[%s] was destroyed but it is pending!', self.coro_id)

    def task_handle(self):
        handle = TaskHandler(close=self.close, closed=lambda: self.task_closed)
        return handle