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
  60     cmd = list[4].split(" ")[0]
  61
  62     if cmd == "mkload":
  63         fp = list[4].split(" ")[1]
  64
  65         try:
  66                 pkfile = open(fp,"r")
  67                 statetab = pickle.load(pkfile)
  68         except IOError:
  69             print "Could not load db: Doesn't exist\n"
  70
  71
  72 #pickles out the state to a file
  73 def markov_dump(c,list,bot):
  74     global statetab
  75     global lock
  76
  77     lock.acquire()
  78
  79     cmd = list[4].split(" ")[0]
  80
  81     if cmd == "mkdump":
  82         fp = list[4].split(" ")[1]
  83
  84         pkfile = open(fp,"w+")
  85
  86         pickle.dump(statetab,pkfile)
  87
  88     lock.release()
  89
  90 def markov_learn(c,list,bot):
  91     """ Should not be called directly """
  92     global lock
  93
  94     lock.acquire()
  95
  96     words = [x for x in list[4].split(" ") if not x.isspace()]
  97
  98     global statetab
  99     global w1
 100     global w2
 101
 102     w1 = w2 = "\n"
 103
 104     #go through every word and put them in a hash table.
 105     #EX the sentence "Mary had a little lamb"
 106     #first iteration, w1 and w2 are both empty.
 107     #statetab[w1][w2] doesn't exist, so make it and set
 108     #statetab[""][""] to Mary.
 109     #
 110     #Then, set w1 to w2 and w2 to i, so the chain moves forward.
 111     for i in words:
 112         statetab.setdefault((w1,w2),[]).append(i)
 113         w1,w2 = w2, i
 114
 115     statetab.setdefault((w1,w2),[]).append("\n")
 116
 117     if nickmatch.search(list[4]) and bot.autorep == 1 and random.randint(0,10) == 0:
 118         tmp = emit_chain(random.choice(list[4].split(" ")))
 119
 120         if len(tmp) <= 2:
 121             return
 122
 123         c.privmsg(list[5], "%s: %s" % (list[0],tmp))
 124         return
 125
 126     #randomly reply
 127     if random.randint(0,15) == 0 and bot.talk == 1:
 128         c.privmsg(list[5], "%s" % (emit_chain(random.choice(list[4].split(" ")))))
 129
 130     lock.release()
 131
 132 def emit_chain(key):
 133     global statetab
 134     global w1
 135     global w2
 136
 137     i = 0
 138
 139     w1 = w2 = "\n"
 140
 141     newword = ""
 142
 143     #make the first word the key if its not a space
 144 #    if(key != " "):
 145 #        retval = key + " "
 146 #    else:
 147     retval = ""
 148
 149     if key != " ":
 150         w2 = key
 151
 152     while 1:
 153         try:
 154             newword = random.choice(statetab[(w1,w2)])
 155         except KeyError:
 156             return retval
 157
 158         retval = retval + newword + " "
 159         w1,w2 = w2,newword
 160
 161         i = i + 1
 162
 163         #max of rand words if we don't hit a space or other error
 164         if i >= random.randint(5,50):
 165             return retval
 166
 167     return retval
 168
 169 def markov_talk(c,list,bot):
 170     """ Makes the markov chain talk to you """
 171     global last
 172
 173     cmd = list[4].split(" ")[0]
 174
 175     try:
 176         key = list[4].split(" ")[1]
 177     except IndexError:
 178         key = " "
 179
 180     if list[3] and cmd == "talk":
 181         tmp = emit_chain(key);
 182         if len(tmp) <= 2:
 183             return
 184         last = tmp
 185         c.privmsg(list[5],"%s" % tmp)
 186
 187
 188
 189 def tweet(c, args, bot):
 190         cmd = args[4].split(" ")[0]
 191
 192         if cmd == "tweet":
 193                 api.PostUpdate(last)
 194                 c.privmsg(args[5], "Updated Twitter with message: %s" % last)