diff --git a/command.py b/command.py index e429d4fe..25d803b0 100644 --- a/command.py +++ b/command.py @@ -55,7 +55,7 @@ class CommandRegistry: @functools.wraps(func) def wrapper(*args, **kwargs): - func(*args, **kwargs) + return func(*args, **kwargs) return wrapper diff --git a/commands/note.py b/commands/note.py index d40855f3..a90e6f8a 100644 --- a/commands/note.py +++ b/commands/note.py @@ -5,7 +5,7 @@ import pytz from command import CommandRegistry from commands import core -from interactive import get_session, has_session, remove_session +from interactive import * from little_shit import get_default_db_path, get_source, get_target __registry__ = cr = CommandRegistry() @@ -32,9 +32,9 @@ _cmd_remove = 'note.remove' @cr.register('记笔记', '添加笔记') @cr.register('take', 'add', hidden=True) @cr.restrict(group_admin_only=True) -def take(args_text, ctx_msg, force=False): +def take(args_text, ctx_msg, allow_interactive=True): source = get_source(ctx_msg) - if not force and (not args_text or has_session(source, _cmd_take)): + if allow_interactive and (not args_text or has_session(source, _cmd_take)): # Be interactive return _take_interactively(args_text, ctx_msg, source) @@ -74,9 +74,9 @@ def list_all(_, ctx_msg): @cr.register('删除笔记') @cr.register('remove', 'delete', hidden=True) @cr.restrict(group_admin_only=True) -def remove(args_text, ctx_msg, force=False): +def remove(args_text, ctx_msg, allow_interactive=True): source = get_source(ctx_msg) - if not force and (not args_text or has_session(source, _cmd_remove)): + if allow_interactive and (not args_text or has_session(source, _cmd_remove)): # Be interactive return _remove_interactively(args_text, ctx_msg, source) @@ -123,7 +123,7 @@ def _take_interactively(args_text, ctx_msg, source): s.state += 1 def save_content(s, a, c): - take(a, c, force=True) + take(a, c, allow_interactive=False) return True if _cmd_take not in _state_machines: @@ -144,7 +144,7 @@ def _remove_interactively(args_text, ctx_msg, source): s.state += 1 def remove_note(s, a, c): - remove(a, c, force=True) + remove(a, c, allow_interactive=False) return True if _cmd_remove not in _state_machines: diff --git a/commands/scheduler.py b/commands/scheduler.py index d36d007b..c671272f 100644 --- a/commands/scheduler.py +++ b/commands/scheduler.py @@ -20,7 +20,7 @@ _scheduler = BackgroundScheduler( 'default': SQLAlchemyJobStore(url=_db_url) }, executors={ - 'default': ProcessPoolExecutor(max_workers=2) + 'default': ProcessPoolExecutor(max_workers=5) }, timezone=pytz.timezone('Asia/Shanghai') ) @@ -152,16 +152,44 @@ def add_job(args_text, ctx_msg, internal=False): @cr.restrict(full_command_only=True, group_admin_only=True) def remove_job(args_text, ctx_msg, internal=False): job_id_without_suffix = args_text.strip() + if not job_id_without_suffix: + _send_text('请指定计划任务的 ID', ctx_msg, internal) + return False job_id = job_id_without_suffix + '_' + get_target(ctx_msg) try: _scheduler.remove_job(job_id, 'default') _send_text('成功删除计划任务 ' + job_id_without_suffix, ctx_msg, internal) + return True except JobLookupError: _send_text('没有找到计划任务 ' + job_id_without_suffix, ctx_msg, internal) + return False + + +@cr.register('get_job', 'get-job', 'get') +@cr.restrict(full_command_only=True) +def get_job(args_text, ctx_msg, internal=False): + job_id_without_suffix = args_text.strip() + if not job_id_without_suffix: + _send_text('请指定计划任务的 ID', ctx_msg, internal) + return None + job_id = job_id_without_suffix + '_' + get_target(ctx_msg) + job = _scheduler.get_job(job_id, 'default') + if internal: + return job + reply = '找到计划任务如下:\n' + reply += 'ID:' + job_id_without_suffix + '\n' + reply += '下次触发时间:\n%s\n' % job.next_run_time.strftime('%Y-%m-%d %H:%M') + reply += '命令:\n' + command_list = job.kwargs['command_list'] + if len(command_list) > 1: + reply += reduce(lambda x, y: x[0] + ' ' + x[1] + '\n' + y[0] + ' ' + y[1], command_list) + else: + reply += command_list[0][0] + ' ' + command_list[0][1] + _send_text(reply, ctx_msg, internal) @cr.register('list_jobs', 'list-jobs', 'list') -@cr.restrict(full_command_only=True, group_admin_only=True) +@cr.restrict(full_command_only=True) def list_jobs(_, ctx_msg, internal=False): target = get_target(ctx_msg) job_id_suffix = '_' + target @@ -173,6 +201,7 @@ def list_jobs(_, ctx_msg, internal=False): job_id = job.id[:-len(job_id_suffix)] command_list = job.kwargs['command_list'] reply = 'ID:' + job_id + '\n' + reply += '下次触发时间:\n%s\n' % job.next_run_time.strftime('%Y-%m-%d %H:%M') reply += '命令:\n' if len(command_list) > 1: reply += reduce(lambda x, y: x[0] + ' ' + x[1] + '\n' + y[0] + ' ' + y[1], command_list) @@ -213,7 +242,7 @@ def _send_add_job_help_msg(ctx_msg, internal): ) core.echo( '--multi 为可选项,表示读取多条命令\n' - 'job_id 为必填项,允许使用符合正则 [_\-a-zA-Z0-9] 的字符,作为计划任务的唯一标识\n' + 'job_id 为必填项,允许使用符合正则 [_\-a-zA-Z0-9] 的字符,作为计划任务的唯一标识,如果指定重复的 ID,则会覆盖原先已有的\n' 'command 为必填项,从 job_id 之后第一个非空白字符开始,如果加了 --multi 选项,则每行算一条命令,否则一直到消息结束算作一整条命令(注意这里的命令不要加 / 前缀)\n' '\n' '例 1:\n' diff --git a/commands/zhihu.py b/commands/zhihu.py index 0229f891..9b8a365b 100644 --- a/commands/zhihu.py +++ b/commands/zhihu.py @@ -3,30 +3,35 @@ from datetime import date, timedelta import requests -from little_shit import SkipException from command import CommandRegistry from commands import core +from commands import scheduler +from interactive import * +from little_shit import SkipException, get_source __registry__ = cr = CommandRegistry() -@cr.register('zhihu', 'zhihu-daily', '知乎日报') +@cr.register('zhihu-daily', 'zhihu', '知乎日报') def zhihu_daily(args_text, ctx_msg): - param = args_text.strip() + arg = args_text.strip() reply = None try: - if not param: + if not arg: sub_url = '/latest' - elif re.match('\d{8}', param) and param >= '20130519': - thedate = date(year=int(param[:4]), month=int(param[4:6]), day=int(param[6:])) - sub_url = '/before/' + (thedate + timedelta(days=1)).strftime('%Y%m%d') else: - reply = '命令格式错误,正确的命令格式:\n' \ - '/zhihu\n' \ - '或\n' \ - '/zhihu 20161129\n' \ - '注意如果指定日期,格式一定要对,且日期需在 20130519 之后。' - raise SkipException + m = re.match('(\d{4})-(\d{2})-(\d{2})', arg) + # and arg >= '20130519': + if m and ''.join(m.groups()) >= '20130519': + thedate = date(year=int(m.group(1)), month=int(m.group(2)), day=int(m.group(3))) + sub_url = '/before/' + (thedate + timedelta(days=1)).strftime('%Y%m%d') + else: + reply = '命令格式错误,正确的命令格式:\n' \ + '/zhihu\n' \ + '或\n' \ + '/zhihu 2016-11-29\n' \ + '注意如果指定日期,格式一定要对,且日期需在 2013-05-19 之后(这一天知乎日报诞生)。' + raise SkipException full_url = 'https://news-at.zhihu.com/api/4/news' + sub_url resp = requests.get( full_url, @@ -58,5 +63,108 @@ def zhihu_daily(args_text, ctx_msg): raise SkipException except SkipException: reply = reply if reply else '发生了未知错误……' - pass core.echo(reply, ctx_msg) + + +_cmd_subscribe = 'zhihu.subscribe' +_scheduler_job_id = _cmd_subscribe + + +@cr.register('订阅知乎日报') +@cr.register('subscribe', hidden=True) +@cr.restrict(group_admin_only=True) +def subscribe(args_text, ctx_msg, allow_interactive=True): + arg = args_text.strip() + source = get_source(ctx_msg) + if allow_interactive and (not arg or has_session(source, _cmd_subscribe)): + # Be interactive + return _subscribe_interactively(args_text, ctx_msg, source) + + force = False + if arg.startswith('-f '): + force = True + arg = arg.split(' ', 1)[1].strip() + reply = None + try: + m = re.match('([0-1]\d|[2][0-3])(?::|:)?([0-5]\d)', arg) + if m: + job = scheduler.get_job(_scheduler_job_id, ctx_msg, internal=True) + if job and not force: + reply = '已经订阅过了哦~\n' \ + + '下次推送时间:\n' \ + + job.next_run_time.strftime('%Y-%m-%d %H:%M') + '\n' \ + + '如果需要更改推送时间,请先取消订阅再重新订阅,' \ + + '或在订阅命令的时间参数前面加 -f 来强制更新推送时间' + raise SkipException + job = scheduler.add_job( + '-M %s -H %s %s zhihu.zhihu-daily' % (m.group(2), m.group(1), _scheduler_job_id), + ctx_msg, + internal=True + ) + if job: + # Succeeded to add a job + reply = '订阅成功,我会在每天 %s 推送哦~' % ':'.join((m.group(1), m.group(2))) + else: + reply = '订阅失败,可能后台出了点问题呢~' + else: + reply = '命令格式错误,正确的命令格式:\n' \ + '/zhihu.subscribe\n' \ + '或\n' \ + '/zhihu.subscribe [-f] 20:30\n' + except SkipException: + reply = reply if reply else '发生了未知错误……' + core.echo(reply, ctx_msg) + + +@cr.register('取消订阅知乎日报') +@cr.register('unsubscribe', hidden=True) +@cr.restrict(group_admin_only=True) +def unsubscribe(_, ctx_msg): + if scheduler.remove_job(_scheduler_job_id, ctx_msg, internal=True): + core.echo('取消订阅成功~', ctx_msg) + else: + core.echo('还没有订阅过哦~', ctx_msg) + + +_state_machines = {} + + +def _subscribe_interactively(args_text, ctx_msg, source): + def confirm_override(s, a, c): + job = scheduler.get_job(_scheduler_job_id, c, internal=True) + if job: + core.echo('先前已经订阅过了哦~\n' + + '下次推送时间:\n' + + job.next_run_time.strftime('%Y-%m-%d %H:%M') + '\n' + + '要更改推送时间吗?\n' + + '回复 1 继续,回复 0 放弃', c) + s.data['need_confirm'] = True + else: + s.data['need_confirm'] = False + wait_for_time(s, a, c) + s.state += 1 + + def wait_for_time(s, a, c): + if s.data['need_confirm']: + if a.strip() != '1': + # Cancel + core.echo('已放弃更改~') + return True + core.echo('请发送想要获取推送的时间(格式如 20:05):', c) + s.state += 1 + + def save(s, a, c): + subscribe('-f ' + a, c, allow_interactive=False) + return True + + if _cmd_subscribe not in _state_machines: + _state_machines[_cmd_subscribe] = ( + confirm_override, # 0 + wait_for_time, # 1 + save # 2 + ) + + sess = get_session(source, _cmd_subscribe) + if _state_machines[_cmd_subscribe][sess.state](sess, args_text, ctx_msg): + # Done + remove_session(source, _cmd_subscribe) diff --git a/config.py b/config.py index 7cafd0c4..c8485c0b 100644 --- a/config.py +++ b/config.py @@ -1,6 +1,6 @@ config = { 'fallback_command': 'core.chat', - 'command_start_flags': ('/', '/'), + 'command_start_flags': ('/', '/', '来,'), 'command_name_separators': ('\.', '->', '::', '/'), # Regex 'command_args_start_flags': (',', ':', ', ', ': '), # Regex }