prime_decomposition_bot.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import sys
import time
import json
import time
import threading
import datetime
import Queue
import nzmath.factor.methods as methods

from slackclient import SlackClient

TOKEN        = "xoxb-00000000000-XXXXXXXXXXXXXXXXXXXXXXXX" # Your API token
CHANNEL      = "random"

CALC_TIMEOUT = 60
MAX_NUMBER   = 10**60
WAIT_TIME    = 1
NUM_WORKERS  = 4
NUM_TRY      = 10

sc = None
q = Queue.Queue()

def factor(number):
    result = Queue.Queue()
    def target(number, result):
        result.put(methods.factor(number))
    t = threading.Thread(target=target, args=(number, result))
    t.start()
    t.join(CALC_TIMEOUT)
    return [] if result.empty() else result.get()

def worker():
    global sc
    
    while True:
        number = q.get()
        if number < 2 or not sc:
            continue
        if number > MAX_NUMBER:
            text = u"%d は大きすぎるので、もう少し小さい数にしてもらえませんか?" % number
            sc.rtm_send_message(CHANNEL, text)
            continue
        factors = [n for (n, m) in factor(number) for i in range(m)]
        text = None
        if len(factors) == 0:
            text = u"頑張って計算したんですけど %d を素因数分解できませんでした… :pensive:" % number
        elif len(factors) == 1:
            text = u"%d は素数です!" % number
        else:
            text = u"%d は %s に素因数分解されます!" % (number, u" と ".join(map(str, factors)))
        if text:
            sc.rtm_send_message(CHANNEL, text)

def main(args):
    global sc
    
    for i in range(NUM_WORKERS):
        t = threading.Thread(target=worker)
        t.daemon = True
        t.start()

    for n in range(NUM_TRY):
        sc = SlackClient(TOKEN)
        if sc.rtm_connect():
            while True:
                try:
                    records = sc.rtm_read()
                except:
                    print "接続が切断されました。再接続します。試行回数: %d" % (n + 1)
                    break
                for record in records:
                    if "text" in record:
                        text = record["text"]
    
                        ### for log
                        print datetime.datetime.fromtimestamp(float(record["ts"])).strftime("%Y-%m-%dT%H:%M:%S") if "ts" in record else "_" * 17,
                        print record["user"] if "user" in record else "_________",
                        print text.encode("utf-8")
                        sys.stdout.flush()
    
                        if text.isdigit():
                            try:
                                q.put(int(text))
                            except:
                                pass
                time.sleep(WAIT_TIME)
        else:
            print "接続に失敗しました。TOKENが間違っていませんか?"
            time.sleep(60)

if __name__ == "__main__":
    main(sys.argv)