kemonine
/
mastodon_bot
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity
Browse Source

botstreamlisteners.py cleanup

merge-requests/1/head
Kirk Strauser 5 years ago
parent
90757d777a
commit
1d596fdf0f
2 changed files with 233 additions and 221 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 452
      botstreamlisteners.py
  2. 2
      toot_bot.py

452
BotStreamListeners.py → botstreamlisteners.py
Unescape View File

@ -1,220 +1,232 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Adjust PyLint settings
# pylint: disable=C0301
''' Different stream listener implementations used by bots '''
import sys, pprint
import sqlite3
import copy
from mastodon import Mastodon, StreamListener
from Toot import toot
class WelcomeBot(StreamListener):
''' Implementation of the Mastodon.py StreamListener class for welcome bot purposes '''
def __init__(self, bot_config):
StreamListener.__init__(self)
self.bot_config = bot_config
# Get access to cache
self.conn = sqlite3.connect(self.bot_config['welcome']['cache_file'])
self.cursor = self.conn.cursor()
# Ensure cache table has been created
self.cursor.execute('''create table if not exists welcome_cache (
username varchar(2048) primary key,
seen_timestamp timestamp
);'''
)
self.cursor.execute('''create table if not exists toot_cache (
toot_id integer primary key
);'''
)
self.conn.commit()
# Replay any toots that were missed while offline and welcome new users
self.replay_toots(True)
def __del__(self):
# Cleanup connection to sqlite database for welcome cache
self.conn.commit()
self.conn.close()
def fetch_remaining(self, mastodon, first_page):
''' Work around for odd behavior in Mastodon.py's official fetch_remaining code '''
# FIXME: Remove this method when below GitHub issue is closed and a new release available
# FIXME: Don't forget to update minimum version of Mastodon.py to match when the fix is released
# https://github.com/halcy/Mastodon.py/issues/59
first_page = copy.deepcopy(first_page)
all_pages = []
current_page = first_page
while current_page != None and current_page:
all_pages.extend(current_page)
current_page = mastodon.fetch_next(current_page)
return all_pages
def replay_toots(self, recurse=False):
''' Replay toots that were posted while the bot was offline '''
# Setup Mastodon API
mastodon = Mastodon(
client_id=self.bot_config['config']['client_cred_file'],
access_token=self.bot_config['config']['user_cred_file'],
api_base_url=self.bot_config['config']['api_base_url']
)
self.cursor.execute('select max(toot_id) from toot_cache;')
last_seen_toot_id = self.cursor.fetchone()[0]
if last_seen_toot_id is None:
last_seen_toot_id = 0
first_page = mastodon.timeline_local(since_id=last_seen_toot_id)
all_pages = self.fetch_remaining(mastodon, first_page)
# Catch up ALL welcome messages that may have been missed
for status in all_pages:
self.welcome_user(status)
# Update max seen toot id (any calls to welcome_user will move the max)
self.cursor.execute('select max(toot_id) from toot_cache;')
last_seen_toot_id = self.cursor.fetchone()[0]
# Update last seen toot id
new_last_seen_toot_id = -1
if all_pages:
new_last_seen_toot_id = all_pages[0]['id']
if new_last_seen_toot_id > last_seen_toot_id:
self.cursor.execute('insert into toot_cache values (?)', (new_last_seen_toot_id,))
self.conn.commit()
# Recurse in case the catch up took long enough for more toots to enter the public timeline
# Do this only once to be safe
if recurse:
self.replay_toots() # Recurse ONCE to catch up on any missing toots posted while doing initial catch up
def welcome_user(self, status):
''' Method that sets up toot and welcomes new users (method due to use in multiple places) '''
toot_id = status['id']
federated = '@' in status['account']['acct']
username = status['account']['acct']
timestamp = status['created_at']
visibility = status['visibility']
# Cache toot
self.cursor.execute('insert into toot_cache values (?)', (toot_id,))
self.conn.commit()
# Welcome any user who's posted publicly
if visibility == 'public' and not federated:
# Check if username has been seen for welcome
self.cursor.execute('select count(1) as found from welcome_cache where username = ?', (username,))
if self.cursor.fetchone()[0] > 0:
return
# Send welcome toot
toot(self.bot_config, username=username)
# Cache user to avoid duping welcome messages
self.cursor.execute('insert into welcome_cache values (?, ?)', (username, timestamp))
self.conn.commit()
def on_update(self, status):
'''A new status has appeared! 'status' is the parsed JSON dictionary
describing the status.'''
self.welcome_user(status)
def on_notification(self, notification):
'''A new notification. 'notification' is the parsed JSON dictionary
describing the notification.'''
# We don't care if notifications come through our bot / curation account
# Leave handling notifications/folow up to the admins and e-mail notifications
pass
def on_delete(self, status_id):
'''A status has been deleted. status_id is the status' integer ID.'''
# Remove the status from the toot_cache if we see a delete
self.cursor.execute('delete from toot_cache where toot_id = ?', (status_id,))
self.conn.commit()
def handle_heartbeat(self):
'''The server has sent us a keep-alive message. This callback may be
useful to carry out periodic housekeeping tasks, or just to confirm
that the connection is still open.'''
# Consistently/constantly trim the toot cache to the most recent seen toot
self.cursor.execute('delete from toot_cache where toot_id <= ((select max(toot_id) from toot_cache) - 1);')
self.conn.commit()
class CurationBot(StreamListener):
''' Implementation of the Mastodon.py StreamListener class for curation bot purposes '''
def __init__(self, bot_config):
StreamListener.__init__(self)
self.bot_config = bot_config
# Get access to cache
self.conn = sqlite3.connect(self.bot_config['curation']['cache_file'])
self.cursor = self.conn.cursor()
# Ensure cache table has been created
self.cursor.execute('''create table if not exists toot_cache (
toot_id integer primary key,
federated bool,
username varchar(2048),
is_reply bool,
is_boost bool,
toot_timestamp timestamp,
favorites integer,
boosts integer
);'''
)
self.conn.commit()
def __del__(self):
# Cleanup connection to sqlite database for rss cache
self.conn.close()
def on_update(self, status):
'''A new status has appeared! 'status' is the parsed JSON dictionary
describing the status.'''
if status['visibility'] == 'public':
toot_id = status['id']
federated = '@' in status['account']['acct']
username = status['account']['acct']
is_reply = status['in_reply_to_id'] is not None
is_boost = status['reblog'] is not None
timestamp = status['created_at']
favorites = status['favourites_count']
boosts = status['reblogs_count']
#tags = status['tags'] # TODO: Add support for tags? - Will need many to 1 relationship and a cross table
# Ensure a toot isn't cached twice for some odd reason
self.cursor.execute('select count(1) as found from toot_cache where toot_id = ?', (toot_id,))
if self.cursor.fetchone()[0] > 0:
return
# Cache toot
self.cursor.execute('insert into toot_cache values (?, ?, ?, ?, ?, ?, ?, ?)', (toot_id, federated, username, is_reply, is_boost, timestamp, favorites, boosts))
self.conn.commit()
def on_notification(self, notification):
'''A new notification. 'notification' is the parsed JSON dictionary
describing the notification.'''
# We don't care if notifications come through our bot / curation account
# Leave handling notifications/folow up to the admins and e-mail notifications
pass
def on_delete(self, status_id):
'''A status has been deleted. status_id is the status' integer ID.'''
# Remove the status from the toot_cache if we see a delete
self.cursor.execute('delete from toot_cache where toot_id = ?', (status_id,))
self.conn.commit()
def handle_heartbeat(self):
'''The server has sent us a keep-alive message. This callback may be
useful to carry out periodic housekeeping tasks, or just to confirm
that the connection is still open.'''
print('!!!!!!!!heartbeat!!!!!!!!')
print(' we should probably update statuses and whatnot here or at least do some housekeeping for old toots')
#!/usr/bin/env python3
''' Different stream listener implementations used by bots '''
import copy
import sqlite3
from toot import toot
from mastodon import Mastodon, StreamListener
class WelcomeBot(StreamListener):
''' Implementation of the Mastodon.py StreamListener class for welcome bot purposes '''
def __init__(self, bot_config):
StreamListener.__init__(self)
self.bot_config = bot_config
# Get access to cache
self.conn = sqlite3.connect(self.bot_config['welcome']['cache_file'])
self.cursor = self.conn.cursor()
# Ensure cache table has been created
self.cursor.execute('''\
CREATE TABLE IF NOT EXISTS welcome_cache (
username VARCHAR(2048) PRIMARY KEY,
seen_timestamp TIMESTAMP
);
''')
self.cursor.execute('''\
CREATE TABLE IF NOT EXISTS toot_cache (
toot_id INTEGER PRIMARY KEY
);''')
self.conn.commit()
# Replay any toots that were missed while offline and welcome new users
self.replay_toots(True)
def __del__(self):
# Cleanup connection to sqlite database for welcome cache
self.conn.commit()
self.conn.close()
def fetch_remaining(self, mastodon, first_page):
''' Work around for odd behavior in Mastodon.py's official fetch_remaining code '''
# FIXME: Remove this method when below GitHub issue is closed
# and a new release available FIXME: Don't forget to update
# minimum version of Mastodon.py to match when the fix is
# released https://github.com/halcy/Mastodon.py/issues/59
first_page = copy.deepcopy(first_page)
all_pages = []
current_page = first_page
while current_page is not None and current_page:
all_pages.extend(current_page)
current_page = mastodon.fetch_next(current_page)
return all_pages
def replay_toots(self, recurse=False):
''' Replay toots that were posted while the bot was offline '''
# Setup Mastodon API
mastodon = Mastodon(
client_id=self.bot_config['config']['client_cred_file'],
access_token=self.bot_config['config']['user_cred_file'],
api_base_url=self.bot_config['config']['api_base_url']
)
self.cursor.execute('select max(toot_id) from toot_cache;')
last_seen_toot_id = self.cursor.fetchone()[0]
if last_seen_toot_id is None:
last_seen_toot_id = 0
first_page = mastodon.timeline_local(since_id=last_seen_toot_id)
all_pages = self.fetch_remaining(mastodon, first_page)
# Catch up ALL welcome messages that may have been missed
for status in all_pages:
self.welcome_user(status)
# Update max seen toot id (any calls to welcome_user will move the max)
self.cursor.execute('select max(toot_id) from toot_cache;')
last_seen_toot_id = self.cursor.fetchone()[0]
# Update last seen toot id
new_last_seen_toot_id = -1
if all_pages:
new_last_seen_toot_id = all_pages[0]['id']
if new_last_seen_toot_id > last_seen_toot_id:
self.cursor.execute('insert into toot_cache values (?)', (new_last_seen_toot_id,))
self.conn.commit()
# Recurse in case the catch up took long enough for more toots to enter the public timeline
# Do this only once to be safe
if recurse:
# Recurse ONCE to catch up on any missing toots posted while doing initial catch up
self.replay_toots()
def welcome_user(self, status):
'''Method that sets up toot and welcomes new users (method due to use
in multiple places)
'''
toot_id = status['id']
federated = '@' in status['account']['acct']
username = status['account']['acct']
timestamp = status['created_at']
visibility = status['visibility']
# Cache toot
self.cursor.execute('insert into toot_cache values (?)', (toot_id,))
self.conn.commit()
# Welcome any user who's posted publicly
if visibility == 'public' and not federated:
# Check if username has been seen for welcome
self.cursor.execute('select count(1) as found from welcome_cache where username = ?',
(username,))
if self.cursor.fetchone()[0] > 0:
return
# Send welcome toot
toot(self.bot_config, username=username)
# Cache user to avoid duping welcome messages
self.cursor.execute('insert into welcome_cache values (?, ?)', (username, timestamp))
self.conn.commit()
def on_update(self, status):
'''A new status has appeared! 'status' is the parsed JSON dictionary
describing the status.'''
self.welcome_user(status)
def on_notification(self, notification):
'''A new notification. 'notification' is the parsed JSON dictionary
describing the notification.'''
# We don't care if notifications come through our bot / curation account
# Leave handling notifications/folow up to the admins and e-mail notifications
pass
def on_delete(self, status_id):
'''A status has been deleted. status_id is the status' integer ID.'''
# Remove the status from the toot_cache if we see a delete
self.cursor.execute('delete from toot_cache where toot_id = ?', (status_id,))
self.conn.commit()
def handle_heartbeat(self):
'''The server has sent us a keep-alive message. This callback may be
useful to carry out periodic housekeeping tasks, or just to confirm
that the connection is still open.'''
# Consistently/constantly trim the toot cache to the most recent seen toot
self.cursor.execute('''\
DELETE FROM toot_cache WHERE toot_id <= ((SELECT MAX(toot_id) FROM toot_cache) - 1);
''')
self.conn.commit()
class CurationBot(StreamListener):
''' Implementation of the Mastodon.py StreamListener class for curation bot purposes '''
def __init__(self, bot_config):
StreamListener.__init__(self)
self.bot_config = bot_config
# Get access to cache
self.conn = sqlite3.connect(self.bot_config['curation']['cache_file'])
self.cursor = self.conn.cursor()
# Ensure cache table has been created
self.cursor.execute('''\
CREATE TABLE IF NOT EXISTS toot_cache (
toot_id INTEGER PRIMARY KEY,
federated BOOL,
username VARCHAR(2048),
is_reply BOOL,
is_boost BOOL,
toot_timestamp TIMESTAMP,
favorites INTEGER,
boosts INTEGER
);''')
self.conn.commit()
def __del__(self):
# Cleanup connection to sqlite database for rss cache
self.conn.close()
def on_update(self, status):
'''A new status has appeared! 'status' is the parsed JSON dictionary
describing the status.'''
if status['visibility'] == 'public':
toot_id = status['id']
federated = '@' in status['account']['acct']
username = status['account']['acct']
is_reply = status['in_reply_to_id'] is not None
is_boost = status['reblog'] is not None
timestamp = status['created_at']
favorites = status['favourites_count']
boosts = status['reblogs_count']
# TODO: Add support for tags? - Will need many to 1 relationship and a cross table
# tags = status['tags']
# Ensure a toot isn't cached twice for some odd reason
self.cursor.execute('select count(1) as found from toot_cache where toot_id = ?',
(toot_id,))
if self.cursor.fetchone()[0] > 0:
return
# Cache toot
self.cursor.execute(
'insert into toot_cache values (?, ?, ?, ?, ?, ?, ?, ?)',
(toot_id, federated, username, is_reply, is_boost, timestamp, favorites, boosts)
)
self.conn.commit()
def on_notification(self, notification):
'''A new notification. 'notification' is the parsed JSON dictionary
describing the notification.'''
# We don't care if notifications come through our bot / curation account
# Leave handling notifications/folow up to the admins and e-mail notifications
pass
def on_delete(self, status_id):
'''A status has been deleted. status_id is the status' integer ID.'''
# Remove the status from the toot_cache if we see a delete
self.cursor.execute('delete from toot_cache where toot_id = ?', (status_id,))
self.conn.commit()
def handle_heartbeat(self):
'''The server has sent us a keep-alive message. This callback may be
useful to carry out periodic housekeeping tasks, or just to confirm
that the connection is still open.'''
print('!!!!!!!!heartbeat!!!!!!!!')
print(' we should probably update statuses and whatnot here or at least do some '
'housekeeping for old toots')

2
toot_bot.py
Unescape View File

@ -13,8 +13,8 @@ import feedparser
from mastodon import Mastodon
from ruamel.yaml import YAML
from botstreamlisteners import CurationBot, WelcomeBot
from toot import toot
from BotStreamListeners import CurationBot, WelcomeBot
def init(config):

Write Preview
Loading…
Cancel
Save