From f8e490dc044c3597d7dd6a1a4e083f23758584b1 Mon Sep 17 00:00:00 2001 From: Klaus-Uwe Mitterer Date: Thu, 23 Mar 2017 17:01:05 +0100 Subject: [PATCH] Add mention streaming --- bottools/methods.py | 31 ++++++++++++++++++++++++------- bottools/strings.py | 3 +++ dbtools/__init__.py | 20 ++++++++++++++++++++ setup.py | 2 +- telegrambot.py | 3 ++- twitools/streaming.py | 5 +++-- 6 files changed, 53 insertions(+), 11 deletions(-) diff --git a/bottools/methods.py b/bottools/methods.py index a6d2f6a..19e61d9 100644 --- a/bottools/methods.py +++ b/bottools/methods.py @@ -221,16 +221,32 @@ def timeline(bot, update, args = [10]): # Streaming mentionstreams = {} +def makeStream(bot, cid): + two = twitools.twoBotHelper(cid) + stream = tweepy.Stream(auth = two.auth, listener = twitools.streaming.BotStreamListener(bot, cid)) + stream.filter(track=[two.whoami()], async=True) + return stream + +try: + for u in dbtools.dbHelper().mentionsOn(): + mentionstreams[u] = makeStream(telegram.Bot(token=setuptools.token()), u) +except Exception as e: + print(e) + def mentionstream(bot, update): global mentionstreams - if update.message.chat_id in mentionstreams: - mentionstreams.pop(update.message.chat_id).disconnect() - else: - two = getTwo(update.message) - stream = tweepy.Stream(auth = two.auth, listener = twitools.streaming.BotStreamListener(bot, update.message.chat_id)) - stream.filter(track=[two.whoami()], async = True) - mentionstreams[update.message.chat_id] = stream + try: + dbtools.dbHelper().toggleMentions(update.message.chat_id) + + if update.message.chat_id in mentionstreams: + mentionstreams.pop(update.message.chat_id).disconnect() + update.message.reply_text(bottools.strings.toggleMentions % bottools.strings.toggleTweetOff) + else: + mentionstreams[update.message.chat_id] = makeStream(bot, update.message.chat_id) + update.message.reply_text(bottools.strings.toggleMentions % bottools.strings.toggleTweetOn) + except tweepy.error.TweepError as e: + bottools.methods.twoExceptions(e, update.message) # Administrator @@ -247,3 +263,4 @@ def restart(bot, update): os.execl(sys.executable, sys.executable, *sys.argv) else: bottools.methods.unknown(bot, update) + diff --git a/bottools/strings.py b/bottools/strings.py index 7c7e369..c7b6d18 100644 --- a/bottools/strings.py +++ b/bottools/strings.py @@ -51,6 +51,9 @@ Have fun!''' toggleTweet = '''Automatic tweeting is now %s.''' +toggleMentions = '''Mention notifications are now %s.''' + + unauth = '''You're leaving already? :( I hope you had a good time with me. If there is anything you would like to tell me or my developers, please drop us a note at: diff --git a/dbtools/__init__.py b/dbtools/__init__.py index ee4508d..8663904 100644 --- a/dbtools/__init__.py +++ b/dbtools/__init__.py @@ -146,6 +146,26 @@ class dbObject: return self.getTStatus(cid) + def getMStatus(self, cid): + try: + self.executeQuery("SELECT mentions FROM tokens WHERE cid = %i;" % int(cid)) + + return True if int(self.cur.fetchone()[0]) == 1 else False + + except: + raise ValueError("No such user: %i" % int(cid)) + + def toggleMentions(self, cid): + self.executeQuery("UPDATE tokens SET mentions = NOT mentions WHERE cid = %i;" % int(cid)) + self.commit() + + return self.getMStatus(cid) + + def mentionsOn(self): + self.executeQuery("SELECT cid FROM tokens WHERE mentions;") + for u in self.getAll(): + yield u[0] + def addFish(self, cid): self.executeQuery("UPDATE tokens SET fish = fish + 1 WHERE cid = %i;" % int(cid)) self.commit() diff --git a/setup.py b/setup.py index 3c98c65..33d02c6 100755 --- a/setup.py +++ b/setup.py @@ -65,7 +65,7 @@ if not db.isInitialized(): db.executeQuery("CREATE TABLE names(`id` TEXT NOT NULL, `name` TEXT NOT NULL, `since` INTEGER NOT NULL, `until` INTEGER, PRIMARY KEY(id, until));") db.executeQuery("CREATE TABLE retweets(id INT PRIMARY KEY, author VARCHAR(30), created_at VARCHAR(30), text TEXT);") db.executeQuery("CREATE TABLE lyrics(id INTEGER PRIMARY KEY AUTOINCREMENT, text VARCHAR(140) NOT NULL, ref INT NOT NULL default '0', tweet_id INT, active BOOLEAN default '0');") - db.executeQuery("CREATE TABLE tokens(`cid` INT PRIMARY KEY, `ato` TEXT, `ase` TEXT, `tweet` BOOLEAN DEFAULT 1, `fish` INT DEFAULT 0);") + db.executeQuery("CREATE TABLE tokens(`cid` INT PRIMARY KEY, `ato` TEXT, `ase` TEXT, `tweet` BOOLEAN DEFAULT 1, `fish` INT DEFAULT 0, `mentions` BOOLEAN DEFAULT 0);") db.executeQuery("CREATE TABLE timelines(`cid` INT, `nr` INT, `tid` INT, PRIMARY KEY(cid, nr));") db.commit() diff --git a/telegrambot.py b/telegrambot.py index 148d10b..427321b 100755 --- a/telegrambot.py +++ b/telegrambot.py @@ -18,13 +18,14 @@ if __name__ == "__main__": updater.dispatcher.add_handler(telegram.ext.CommandHandler("like", bottools.methods.like, pass_args=True)) updater.dispatcher.add_handler(telegram.ext.CommandHandler("login", bottools.methods.auth)) updater.dispatcher.add_handler(telegram.ext.CommandHandler("logout", bottools.methods.unauth)) - updater.dispatcher.add_handler(telegram.ext.CommandHandler("mentionstream", bottools.methods. mentionstream)) + updater.dispatcher.add_handler(telegram.ext.CommandHandler("mentionstream", bottools.methods.mentionstream)) updater.dispatcher.add_handler(telegram.ext.CommandHandler("reply", bottools.methods.reply, pass_args=True)) updater.dispatcher.add_handler(telegram.ext.CommandHandler("restart", bottools.methods.restart)) updater.dispatcher.add_handler(telegram.ext.CommandHandler("retweet", bottools.methods.retweet, pass_args=True)) updater.dispatcher.add_handler(telegram.ext.CommandHandler("start", bottools.methods.start)) updater.dispatcher.add_handler(telegram.ext.CommandHandler("timeline", bottools.methods.timeline, pass_args=True)) updater.dispatcher.add_handler(telegram.ext.CommandHandler("toggletweet", bottools.methods.toggleTweet)) + updater.dispatcher.add_handler(telegram.ext.CommandHandler("togglementions", bottools.methods.mentionstream)) updater.dispatcher.add_handler(telegram.ext.CommandHandler("tweet", bottools.methods.explicitTweet, pass_args=True)) updater.dispatcher.add_handler(telegram.ext.CommandHandler("unauth", bottools.methods.unauth)) updater.dispatcher.add_handler(telegram.ext.CommandHandler("unfollow", bottools.methods.unfollow, pass_args=True)) diff --git a/twitools/streaming.py b/twitools/streaming.py index 1d01bbe..b22324c 100644 --- a/twitools/streaming.py +++ b/twitools/streaming.py @@ -1,7 +1,8 @@ -import dbtools, tweepy +import dbtools, html, tweepy class BotStreamListener(tweepy.StreamListener): - def __init__(self, bot, cid): + def __init__(self, bot, cid, *args, **kwargs): + tweepy.StreamListener.__init__(self, *args, **kwargs) self.bot = bot self.cid = cid