class Bot(object):
"""Slackchat bot for imitating users"""
def __init__(self, token):
# hardcoded
self.blacklist = set(["markovbot", "kitty"])
# hardcoded
self.save_file = ".messagestore"
# Keyed by user_id
self.markovs = defaultdict(Markov)
# Flag to inform polling thread of shutdown
self.running = False
self.slack = Slack(token)
self.thread = None
# Load saved MCs
if os.path.isfile(self.save_file):
with open(self.save_file) as handle:
self.markovs = pickle.load(handle)
# Clean up old versions
for key in self.markovs:
if not self.markovs[key]:
del self.markovs[key]
# Reverse dict lookup
id_list = [k for k in self.markovs if self.slack.users[k].name == "markovbot"]
if id_list:
print "Deleting introspective data"
del self.markovs[id_list[0]]
else:
self.messages = defaultdict(list)
def start(self):
"""Begin polling for messages"""
self.running = True
self.thread = threading.Thread(target=self.poll)
self.thread.start()
print "Started polling thread"
def kill(self):
"""Kill the polling thread"""
self.running = False
self.thread.join()
self.save()
def get_history(self):
"""Download as much history as possible"""
count = 0
for user, text in self.slack.get_history():
if self.add_message(user, text):
count += 1
print "Loaded %d messages from history" % count
def save(self):
"""Save dictionary of MCs
#TODO: Incremental save
"""
with open(self.save_file, 'w+') as handle:
pickle.dump(self.markovs, handle)
def add_message(self, user, text):
"""Add transition counts from a received message"""
if user and self.slack.users[user].name not in self.blacklist:
self.markovs[user].add(text)
return True
else:
return False
def poll(self):
"""Poll slack API for messages
TODO: Break saving state out into separate thread or at least method
"""
count = 0
while self.running:
user, message = self.slack.get_message()
# This is here to make it so we can exit cleanly
# TODO: Something less clunky
if message is None:
continue
count += 1
# hardcoded
if count == 50:
self.save()
count = 0
print "%s: %s" % (user, message)
self.add_message(user, message)
def send_message(self):
"""Choose a user, generate a message for them, and send it
TODO: Break out user selection, and make clever
TODO: Somehow seed by previous message
"""
if self.markovs:
user_id = random.choice([k for k in self.markovs])
text = self.markovs[user_id].generate()
# hardcoded
self.slack.send_message("bot-playground", user_id, text)
else:
# No data loaded
pass
