73 lines
2.5 KiB
Python
73 lines
2.5 KiB
Python
import os
|
|
import importlib
|
|
import aiosqlite
|
|
|
|
import nerimity
|
|
import pylast
|
|
|
|
class Bot(nerimity.Client):
|
|
"""Extended client class for extra functionality."""
|
|
def __init__(self, lastfm_api_key: str, lastfm_api_secret: str, owner_id: int = None, db_file: str = "database.db", *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
|
|
self.owner_id = owner_id
|
|
|
|
# initialize last.fm client
|
|
self.lastfm = pylast.LastFMNetwork(
|
|
api_key = lastfm_api_key,
|
|
api_secret = lastfm_api_secret
|
|
)
|
|
|
|
self.db = None
|
|
self.db_file = db_file
|
|
|
|
|
|
def load_commands(self, commands_dir: str):
|
|
for filename in os.listdir(commands_dir):
|
|
if filename.endswith(".py") and not filename.startswith("_"):
|
|
module_name = f"{commands_dir}.{filename[:-3]}"
|
|
module = importlib.import_module(module_name)
|
|
if hasattr(module, "setup"):
|
|
module.setup(self)
|
|
print(f"Loaded {module_name}")
|
|
print("Registered commands:", self.commands)
|
|
|
|
def is_owner(self, user: nerimity.Member = None):
|
|
if not user: return None
|
|
else:
|
|
if user.id == self.owner_id: return True
|
|
else: return False
|
|
|
|
def _process_slash_commands(self, message: nerimity.Message):
|
|
if message.author.id == self.account.id: return
|
|
super()._process_slash_commands(message)
|
|
|
|
async def _process_commands(self, message: nerimity.Message):
|
|
if message.author.id == self.account.id: return
|
|
await super()._process_commands(message)
|
|
|
|
async def init_db(self):
|
|
print("Initializing database...")
|
|
self.db = await aiosqlite.connect(self.db_file)
|
|
await self.db.execute("""
|
|
CREATE TABLE IF NOT EXISTS lastfm_users (
|
|
user_id TEXT PRIMARY KEY,
|
|
lastfm_username TEXT NOT NULL
|
|
)
|
|
""")
|
|
await self.db.commit()
|
|
|
|
async def set_lastfm(self, user_id: str, username: str):
|
|
await self.db.execute(
|
|
"INSERT OR REPLACE INTO lastfm_users (user_id, lastfm_username) VALUES (?, ?)",
|
|
(user_id, username)
|
|
)
|
|
await self.db.commit()
|
|
|
|
async def get_lastfm(self, user_id: str):
|
|
async with self.db.execute(
|
|
"SELECT lastfm_username FROM lastfm_users WHERE user_id = ?",
|
|
(user_id,)
|
|
) as cursor:
|
|
row = await cursor.fetchone()
|
|
return row[0] if row else None
|