From de9a4e5b83eb26481f84f2e2bce09727b6850f67 Mon Sep 17 00:00:00 2001
From: Richard Chien <richardchienthebest@gmail.com>
Date: Tue, 24 Jul 2018 23:59:45 +0800
Subject: [PATCH] Add CommandSession.switch()

---
 none/command.py | 58 +++++++++++++++++++++++++++++++++++++++++++++----
 none/message.py | 11 ++++++++--
 2 files changed, 63 insertions(+), 6 deletions(-)

diff --git a/none/command.py b/none/command.py
index ca9154af..09870ea5 100644
--- a/none/command.py
+++ b/none/command.py
@@ -170,7 +170,7 @@ class _FurtherInteractionNeeded(Exception):
 class _FinishException(Exception):
     """
     Raised by session.finish() indicating that the command session
-    should be stop and removed.
+    should be stopped and removed.
     """
 
     def __init__(self, result: bool = True):
@@ -180,6 +180,24 @@ class _FinishException(Exception):
         self.result = result
 
 
+class SwitchException(Exception):
+    """
+    Raised by session.switch() indicating that the command session
+    should be stopped and replaced with a new one (going through
+    handle_message() again).
+
+    Since the new context message will go through handle_message()
+    again, the later function should be notified. So this exception
+    is designed to be propagated to handle_message().
+    """
+
+    def __init__(self, new_ctx_message: Message):
+        """
+        :param new_ctx_message: new message which should be placed in context
+        """
+        self.new_ctx_message = new_ctx_message
+
+
 class CommandSession(BaseSession):
     __slots__ = ('cmd', 'current_key', 'current_arg', 'current_arg_text',
                  'current_arg_images', 'args', '_last_interaction', '_running')
@@ -278,6 +296,26 @@ class CommandSession(BaseSession):
             asyncio.ensure_future(self.send(message))
         raise _FinishException
 
+    # noinspection PyMethodMayBeStatic
+    def switch(self, new_ctx_message: Any) -> None:
+        """
+        Finish the session and switch to a new (fake) message context.
+
+        The user may send another command (or another intention as natural
+        language) when interacting with the current session. In this case,
+        the session may not understand what the user is saying, so it
+        should call this method and pass in that message, then NoneBot will
+        handle the situation properly.
+        """
+        if self.is_first_run:
+            # if calling this method during first run,
+            # we think the command is not handled
+            raise _FinishException(result=False)
+
+        if not isinstance(new_ctx_message, Message):
+            new_ctx_message = Message(new_ctx_message)
+        raise SwitchException(new_ctx_message)
+
 
 def parse_command(bot: NoneBot,
                   cmd_string: str) -> Tuple[Optional[Command], Optional[str]]:
@@ -433,9 +471,11 @@ async def _real_run_command(session: CommandSession,
                             ctx_id: str,
                             disable_interaction: bool = False,
                             **kwargs) -> bool:
+    session_overridden = False
     if not disable_interaction:
         # override session only when not disabling interaction
         _sessions[ctx_id] = session
+        session_overridden = True
     try:
         logger.debug(f'Running command {session.cmd.name}')
         session.running = True
@@ -450,10 +490,20 @@ async def _real_run_command(session: CommandSession,
                      f'command {session.cmd.name}')
         # return True because this step of the session is successful
         return True
-    except _FinishException as e:
+    except (_FinishException, SwitchException) as e:
         session.running = False
         logger.debug(f'Session of command {session.cmd.name} finished')
-        if not disable_interaction and ctx_id in _sessions:
+        if session_overridden:
             # the command is finished, remove the session
             del _sessions[ctx_id]
-        return e.result
+
+        if isinstance(e, _FinishException):
+            return e.result
+        elif isinstance(e, SwitchException):
+            # we are guaranteed that the session is not first run here
+            if ctx_id in _sessions:
+                # make sure there is no session waiting
+                del _sessions[ctx_id]
+            logger.debug(f'Session of command {session.cmd.name} switching,'
+                         f'new context message: {e.new_ctx_message}')
+            raise e  # this is intended to be propagated to handle_message()
diff --git a/none/message.py b/none/message.py
index 53676f49..525e350a 100644
--- a/none/message.py
+++ b/none/message.py
@@ -3,7 +3,7 @@ from typing import Dict, Any
 from aiocqhttp.message import MessageSegment
 
 from . import NoneBot
-from .command import handle_command
+from .command import handle_command, SwitchException
 from .log import logger
 from .natural_language import handle_natural_language
 
@@ -23,7 +23,14 @@ async def handle_message(bot: NoneBot, ctx: Dict[str, Any]) -> None:
     else:
         ctx['to_me'] = True
 
-    handled = await handle_command(bot, ctx)
+    while True:
+        try:
+            handled = await handle_command(bot, ctx)
+            break
+        except SwitchException as e:
+            # we are sure that there is no session existing now
+            ctx['message'] = e.new_ctx_message
+            ctx['to_me'] = True
     if handled:
         logger.info(f'Message {ctx["message_id"]} is handled as a command')
         return