From 946d0a11cd67b73d6a6e4d261c93a5e74998c294 Mon Sep 17 00:00:00 2001
From: Richard Chien <richardchienthebest@gmail.com>
Date: Thu, 5 Jul 2018 23:11:00 +0800
Subject: [PATCH] Improve command module and add "schedule" command group

---
 none/argparse.py         |  16 +++++
 none/command.py          | 101 +++++++++++++++-----------
 none/natural_language.py |   5 +-
 none/plugins/schedule.py | 151 +++++++++++++++++++++++++++++++++++++++
 4 files changed, 231 insertions(+), 42 deletions(-)
 create mode 100644 none/argparse.py
 create mode 100644 none/plugins/schedule.py

diff --git a/none/argparse.py b/none/argparse.py
new file mode 100644
index 00000000..763094e5
--- /dev/null
+++ b/none/argparse.py
@@ -0,0 +1,16 @@
+from argparse import *
+
+
+class ParserExit(RuntimeError):
+    def __init__(self, status=0, message=None):
+        self.status = status
+        self.message = message
+
+
+class ArgumentParser(ArgumentParser):
+    def _print_message(self, *args, **kwargs):
+        # do nothing
+        pass
+
+    def exit(self, status=0, message=None):
+        raise ParserExit(status=status, message=message)
diff --git a/none/command.py b/none/command.py
index 6929025e..43cc6029 100644
--- a/none/command.py
+++ b/none/command.py
@@ -37,26 +37,38 @@ class Command:
         self.only_to_me = only_to_me
         self.args_parser_func = None
 
-    async def run(self, session, check_perm: bool = True) -> bool:
+    async def run(self, session, *,
+                  check_perm: bool = True,
+                  dry: bool = False) -> bool:
         """
         Run the command in a given session.
 
         :param session: CommandSession object
         :param check_perm: should check permission before running
-        :return: the command is finished
+        :param dry: just check any prerequisite, without actually running
+        :return: the command is finished (or can be run, given dry == True)
         """
-        if check_perm:
-            has_perm = await perm.check_permission(
-                session.bot, session.ctx, self.permission)
-        else:
-            has_perm = True
+        has_perm = await self._check_perm(session) if check_perm else True
         if self.func and has_perm:
+            if dry:
+                return True
             if self.args_parser_func:
                 await self.args_parser_func(session)
             await self.func(session)
             return True
         return False
 
+    async def _check_perm(self, session) -> bool:
+        """
+        Check if the session has sufficient permission to
+        call the command.
+
+        :param session: CommandSession object
+        :return: the session has the permission
+        """
+        return await perm.check_permission(session.bot, session.ctx,
+                                           self.permission)
+
 
 def on_command(name: Union[str, Tuple[str]], *,
                aliases: Iterable = (),
@@ -226,31 +238,25 @@ class CommandSession(BaseSession):
         return self.args.get(key, default)
 
 
-def _new_command_session(bot: NoneBot,
-                         ctx: Dict[str, Any]) -> Optional[CommandSession]:
+def parse_command(bot: NoneBot,
+                  cmd_string: str) -> Tuple[Optional[Command], Optional[str]]:
     """
-    Create a new session for a command.
-
-    This will attempt to parse the current message as a command,
-    and if succeeded, it then create a session for the command and return.
-    If the message is not a valid command, None will be returned.
+    Parse a command string (typically from a message).
 
     :param bot: NoneBot instance
-    :param ctx: message context
-    :return: CommandSession object or None
+    :param cmd_string: command string
+    :return: (Command object, current arg string)
     """
-    msg = str(ctx['message']).lstrip()
-
     matched_start = None
     for start in bot.config.COMMAND_START:
         # loop through COMMAND_START to find the longest matched start
         curr_matched_start = None
         if isinstance(start, type(re.compile(''))):
-            m = start.search(msg)
+            m = start.search(cmd_string)
             if m and m.start(0) == 0:
                 curr_matched_start = m.group(0)
         elif isinstance(start, str):
-            if msg.startswith(start):
+            if cmd_string.startswith(start):
                 curr_matched_start = start
 
         if curr_matched_start is not None and \
@@ -261,13 +267,13 @@ def _new_command_session(bot: NoneBot,
 
     if matched_start is None:
         # it's not a command
-        return None
+        return None, None
 
-    full_command = msg[len(matched_start):].lstrip()
+    full_command = cmd_string[len(matched_start):].lstrip()
 
     if not full_command:
         # command is empty
-        return None
+        return None, None
 
     cmd_name_text, *cmd_remained = full_command.split(maxsplit=1)
     cmd_name = _aliases.get(cmd_name_text)
@@ -291,11 +297,9 @@ def _new_command_session(bot: NoneBot,
 
     cmd = _find_command(cmd_name)
     if not cmd:
-        return None
-    if cmd.only_to_me and not ctx['to_me']:
-        return None
+        return None, None
 
-    return CommandSession(bot, ctx, cmd, current_arg=''.join(cmd_remained))
+    return cmd, ''.join(cmd_remained)
 
 
 async def handle_command(bot: NoneBot, ctx: Dict[str, Any]) -> bool:
@@ -322,48 +326,65 @@ async def handle_command(bot: NoneBot, ctx: Dict[str, Any]) -> bool:
             del _sessions[ctx_id]
             session = None
     if not session:
-        session = _new_command_session(bot, ctx)
-        if not session:
+        cmd, current_arg = parse_command(bot, str(ctx['message']).lstrip())
+        if not cmd or cmd.only_to_me and not ctx['to_me']:
             return False
+        session = CommandSession(bot, ctx, cmd, current_arg=current_arg)
     return await _real_run_command(session, ctx_id, check_perm=check_perm)
 
 
 async def call_command(bot: NoneBot, ctx: Dict[str, Any],
-                       name: Union[str, Tuple[str]],
-                       args: Dict[str, Any]) -> bool:
+                       name: Union[str, Tuple[str]], *,
+                       current_arg: str = '',
+                       args: Optional[Dict[str, Any]] = None,
+                       check_perm: bool = True,
+                       disable_interaction: bool = False) -> bool:
     """
     Call a command internally.
 
     This function is typically called by some other commands
     or "handle_natural_language" when handling NLPResult object.
 
-    Note: After calling this function, any previous command session
-    will be overridden, even if the command being called here does
-    not need further interaction (a.k.a asking the user for more info).
+    Note: If disable_interaction is not True, after calling this function,
+    any previous command session will be overridden, even if the command
+    being called here does not need further interaction (a.k.a asking
+    the user for more info).
 
     :param bot: NoneBot instance
     :param ctx: message context
     :param name: command name
+    :param current_arg: command current argument string
     :param args: command args
+    :param check_perm: should check permission before running command
+    :param disable_interaction: disable the command's further interaction
     :return: the command is successfully called
     """
     cmd = _find_command(name)
     if not cmd:
         return False
-    session = CommandSession(bot, ctx, cmd, args=args)
+    session = CommandSession(bot, ctx, cmd, current_arg=current_arg, args=args)
     return await _real_run_command(session, context_id(session.ctx),
-                                   check_perm=False)
+                                   check_perm=check_perm,
+                                   disable_interaction=disable_interaction)
 
 
 async def _real_run_command(session: CommandSession,
-                            ctx_id: str, **kwargs) -> bool:
-    _sessions[ctx_id] = session
+                            ctx_id: str,
+                            disable_interaction: bool = False,
+                            **kwargs) -> bool:
+    if not disable_interaction:
+        # override session only when not disabling interaction
+        _sessions[ctx_id] = session
     try:
         res = await session.cmd.run(session, **kwargs)
-        # the command is finished, remove the session
-        del _sessions[ctx_id]
+        if not disable_interaction:
+            # the command is finished, remove the session
+            del _sessions[ctx_id]
         return res
     except _FurtherInteractionNeeded:
+        if disable_interaction:
+            # if the command needs further interaction, we view it as failed
+            return False
         session.last_interaction = datetime.now()
         # return True because this step of the session is successful
         return True
diff --git a/none/natural_language.py b/none/natural_language.py
index 957a1582..b35aa26a 100644
--- a/none/natural_language.py
+++ b/none/natural_language.py
@@ -117,6 +117,7 @@ async def handle_natural_language(bot: NoneBot, ctx: Dict[str, Any]) -> bool:
         logger.debug(results)
         if results and results[0].confidence >= 60.0:
             # choose the result with highest confidence
-            return await call_command(bot, ctx,
-                                      results[0].cmd_name, results[0].cmd_args)
+            return await call_command(bot, ctx, results[0].cmd_name,
+                                      args=results[0].cmd_args,
+                                      check_perm=False)
     return False
diff --git a/none/plugins/schedule.py b/none/plugins/schedule.py
new file mode 100644
index 00000000..9d6edd13
--- /dev/null
+++ b/none/plugins/schedule.py
@@ -0,0 +1,151 @@
+import atexit
+import re
+import shlex
+from typing import Iterable, Tuple, Dict, Any
+
+from apscheduler.jobstores.base import ConflictingIdError
+from apscheduler.jobstores.sqlalchemy import SQLAlchemyJobStore
+from apscheduler.schedulers.asyncio import AsyncIOScheduler
+
+from none import get_bot, CommandGroup, CommandSession, permission as perm
+from none.argparse import ArgumentParser, ParserExit
+from none.command import parse_command, call_command
+from none.helpers import context_id, send
+
+sched = CommandGroup('schedule', permission=perm.PRIVATE | perm.GROUP_ADMIN,
+                     only_to_me=False)
+
+_bot = get_bot()
+
+_scheduler = AsyncIOScheduler(
+    jobstores={
+        'default': SQLAlchemyJobStore(
+            url='sqlite:///' + _bot.get_data_file('db', 'schedule.sqlite'))
+    },
+    timezone='Asia/Shanghai'
+)
+
+if not _scheduler.running:
+    _scheduler.start()
+
+
+@atexit.register
+def _():
+    if _scheduler.running:
+        _scheduler.shutdown()
+
+
+async def _schedule_callback(ctx: Dict[str, Any], name: str,
+                             commands: Iterable[Tuple[Tuple[str], str]],
+                             verbose: bool = False):
+    if verbose:
+        await send(_bot, ctx, f'开始执行计划任务 {name}……')
+    for cmd_name, current_arg in commands:
+        await call_command(_bot, ctx, cmd_name,
+                           current_arg=current_arg,
+                           check_perm=True,
+                           disable_interaction=True)
+
+
+@sched.command('add', aliases=('schedule',))
+async def sched_add(session: CommandSession):
+    parser = ArgumentParser('schedule.add')
+    parser.add_argument('-S', '--second')
+    parser.add_argument('-M', '--minute')
+    parser.add_argument('-H', '--hour')
+    parser.add_argument('-d', '--day')
+    parser.add_argument('-m', '--month')
+    parser.add_argument('-w', '--day-of-week')
+    parser.add_argument('-f', '--force', action='store_true', default=False)
+    parser.add_argument('-v', '--verbose', action='store_true', default=False)
+    parser.add_argument('--name', required=True)
+    parser.add_argument('commands', nargs='+')
+
+    argv = session.get_optional('argv')
+    if not argv:
+        await session.send(_sched_add_help)
+        return
+
+    try:
+        args = parser.parse_args(argv)
+    except ParserExit as e:
+        if e.status == 0:
+            # --help
+            await session.send(_sched_add_help)
+        else:
+            await session.send('参数不足或不正确，请使用 --help 参数查询使用帮助')
+        return
+
+    if not re.match(r'[_a-zA-Z][_\w]*', args.name):
+        await session.send(
+            '计划任务名必须仅包含字母、数字、下划线，且以字母或下划线开头')
+        return
+
+    parsed_commands = []
+    invalid_commands = []
+    for cmd_str in args.commands:
+        cmd, current_arg = parse_command(session.bot, cmd_str)
+        if cmd:
+            tmp_session = CommandSession(session.bot, session.ctx, cmd,
+                                         current_arg=current_arg)
+            if await cmd.run(tmp_session, dry=True):
+                parsed_commands.append((cmd.name, current_arg))
+            else:
+                invalid_commands.append(cmd_str)
+    if invalid_commands:
+        invalid_commands_joined = '\r\n'.join(
+            [f'{i+1}: {c}' for i, c in enumerate(invalid_commands)])
+        await session.send(f'计划任务添加失败，'
+                           f'因为下面的 {len(invalid_commands)} 个命令无法被运行'
+                           f'（命令不存在或权限不够）：\r\n\r\n'
+                           f'{invalid_commands_joined}')
+        return
+
+    job_id = f'{context_id(session.ctx)}/job/{args.name}'
+    trigger_args = {k: v for k, v in args.__dict__.items()
+                    if k in {'second', 'minute', 'hour',
+                             'day', 'month', 'day_of_week'}}
+    try:
+        job = _scheduler.add_job(
+            _schedule_callback,
+            trigger='cron', **trigger_args,
+            id=job_id,
+            kwargs={
+                'ctx': session.ctx,
+                'name': args.name,
+                'commands': parsed_commands,
+                'verbose': args.verbose,
+            },
+            replace_existing=args.force,
+            misfire_grace_time=30
+        )
+    except ConflictingIdError:
+        # a job with same name exists
+        await session.send(f'计划任务 {args.name} 已存在，'
+                           f'若要覆盖请使用 --force 参数')
+        return
+
+    await session.send(f'计划任务 {args.name} 添加成功，下次运行时间：'
+                       f'{job.next_run_time.strftime("%Y-%m-%d %H:%M:%S")}')
+
+
+@sched.command('list')
+async def sched_list(session: CommandSession):
+    pass
+
+
+@sched.command('remove')
+async def sched_remove(session: CommandSession):
+    pass
+
+
+@sched_add.args_parser
+@sched_list.args_parser
+@sched_remove.args_parser
+async def _(session: CommandSession):
+    session.args['argv'] = shlex.split(session.current_arg_text)
+
+
+_sched_add_help = r"""
+使用方法：schedule.add
+""".strip()