modules/markov/markov.py

   1 #markov module for scrappy
   2
   3 #the inspiration for this code comes from
   4 #"The Practice of Programming", so some of the variable
   5 #names and such will be similar.
   6
   7 import random
   8 import pickle
   9 import re
  10 import threading
  11
  12
  13 #EXPERIMENTAL
  14 last = ""
  15 import twitter
  16
  17 api = None
  18 twitteruser = 'scrappybot'
  19 twitterpass = 'mark0vb0t'
  20
  21 nickmatch = None
  22 statetab = {}
  23 lock = None
  24 w1 = w2 = "\n"
  25
  26
  27 def init(scrap):
  28     global nickmatch
  29     global lock
  30     global api
  31
  32     lock = threading.Lock()
  33
  34     scrap.register_event("markov", "msg", markov_learn)
  35     scrap.register_event("markov", "msg", markov_talk)
  36     scrap.register_event("markov", "msg", markov_load)
  37     scrap.register_event("markov", "msg", markov_dump)
  38     scrap.register_event("markov", "msg", markov_stats)
  39     scrap.register_event("markov", "msg", tweet)
  40
  41     nickmatch = re.compile(scrap.nickname)
  42
  43     random.seed()
  44
  45     api = twitter.Api(username=twitteruser, password=twitterpass)
  46
  47 def markov_stats(c,list,bot):
  48     global statetab
  49
  50     cmd = list[4].split(" ")[0]
  51
  52     if cmd == "markov_stats":
  53         c.privmsg(list[5], "words: %d" % len(statetab[("\n","\n")]))
  54         c.privmsg(list[5], "chains: %d" % len(statetab.items()))
  55
  56 #loads in a previously pickled saved state
  57 def markov_load(c,list,bot):
  58     global statetab
  59     global lock
  60
  61     cmd = list[4].split(" ")[0]
  62
  63     if cmd == "mkload":
  64         fp = list[4].split(" ")[1]
  65
  66         try:
  67                 lock.acquire()
  68                 pkfile = open(fp,"r")
  69                 statetab = pickle.load(pkfile)
  70                 lock.release()
  71         except IOError:
  72             print "Could not load db: Doesn't exist\n"
  73
  74
  75 #pickles out the state to a file
  76 def markov_dump(c,list,bot):
  77     global statetab
  78     global lock
  79
  80     lock.acquire()
  81
  82     cmd = list[4].split(" ")[0]
  83
  84     if cmd == "mkdump":
  85         fp = list[4].split(" ")[1]
  86
  87         pkfile = open(fp,"w+")
  88
  89         pickle.dump(statetab,pkfile)
  90
  91     lock.release()
  92
  93 def markov_learn(c,list,bot):
  94     """ Should not be called directly """
  95     global lock
  96
  97     cmd = list[4].split(" ")[0]
  98     if cmd == "talk" or cmd == "markov_stats" or cmd == "mkload" or cmd == "mkdump":
  99             return
 100     lock.acquire()
 101
 102     words = [x.lower().strip() for x in list[4].split(" ") if not x.isspace()]
 103
 104     global statetab
 105     global w1
 106     global w2
 107
 108     w1 = w2 = "\n"
 109
 110     #go through every word and put them in a hash table.
 111     #EX the sentence "Mary had a little lamb"
 112     #first iteration, w1 and w2 are both empty.
 113     #statetab[w1][w2] doesn't exist, so make it and set
 114     #statetab[""][""] to Mary.
 115     #
 116     #Then, set w1 to w2 and w2 to i, so the chain moves forward.
 117     for i in words:
 118         statetab.setdefault((w1,w2),[]).append(i)
 119         w1,w2 = w2, i
 120
 121     statetab.setdefault((w1,w2),[]).append("\n")
 122
 123     if nickmatch.search(list[4]) and bot.autorep == 1 and random.randint(0,10) == 0:
 124         tmp = emit_chain(random.choice(list[4].split(" ")))
 125
 126         if len(tmp) <= 2:
 127             return
 128
 129         c.privmsg(list[5], "%s: %s" % (list[0],tmp))
 130         return
 131
 132     #randomly reply
 133     if random.randint(0,15) == 0 and bot.talk == 1:
 134         c.privmsg(list[5], "%s" % (emit_chain(random.choice(list[4].split(" ")))))
 135
 136     lock.release()
 137
 138 def emit_chain(key):
 139     global statetab
 140     global w1
 141     global w2
 142
 143     i = 0
 144
 145     w1 = w2 = "\n"
 146
 147     newword = ""
 148
 149     #make the first word the key if its not a space
 150 #    if(key != " "):
 151 #        retval = key + " "
 152 #    else:
 153     retval = ""
 154
 155     if key != " ":
 156         w2 = key
 157
 158     while 1:
 159         try:
 160             newword = random.choice(statetab[(w1,w2)]).strip()
 161         except KeyError:
 162             return retval
 163
 164         retval = retval + newword + " "
 165         w1,w2 = w2,newword
 166
 167         i = i + 1
 168
 169         #max of rand words if we don't hit a space or other error
 170         if i >= random.randint(5,50):
 171             return retval
 172
 173     return retval
 174
 175 def markov_talk(c,list,bot):
 176     """ Makes the markov chain talk to you """
 177     global last
 178
 179     cmd = list[4].split(" ")[0]
 180
 181     try:
 182         key = list[4].split(" ")[1].lower()
 183     except IndexError:
 184         key = " "
 185
 186     if list[3] and cmd == "talk":
 187         tmp = emit_chain(key);
 188         if len(tmp) <= 2:
 189             return
 190         last = tmp
 191         c.privmsg(list[5],"%s %s" % (key, tmp))
 192
 193
 194
 195 def tweet(c, args, bot):
 196         cmd = args[4].split(" ")[0]
 197
 198         if cmd == "tweet":
 199                 api.PostUpdate(last)
 200                 c.privmsg(args[5], "Updated Twitter with message: %s" % last)