Update a lot of things

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)