diff --git a/.gitignore b/.gitignore index b16969b..f14b444 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,4 @@ venv *.pickle *.json *.swp +sync_token diff --git a/main.py b/main.py index f0e7245..f35e39d 100755 --- a/main.py +++ b/main.py @@ -1,17 +1,18 @@ #!/usr/bin/env python3 """parkerbot: Matrix bot to generate YouTube (music) playlists from links sent to a channel.""" +import asyncio import os +import pickle import re import sqlite3 -import asyncio -import pickle from datetime import datetime, timedelta + from dotenv import load_dotenv +from google.auth.transport.requests import Request from google_auth_oauthlib.flow import InstalledAppFlow from googleapiclient.discovery import build -from google.auth.transport.requests import Request -from nio import AsyncClient, MatrixRoom, RoomMessageText +from nio import AsyncClient, ClientConfig, RoomMessageText, SyncResponse load_dotenv() DB_PATH = os.getenv("DB_PATH") @@ -28,7 +29,8 @@ conn = sqlite3.connect(DB_PATH) cursor = conn.cursor() -def create_tables(): +def define_tables(): + """Define tables for use with program.""" with conn: cursor.execute( """CREATE TABLE IF NOT EXISTS messages ( @@ -57,6 +59,7 @@ def create_tables(): def get_authenticated_service(): + """Get an authentivated YouTube service.""" credentials = None # The file token.pickle stores the user's access and refresh tokens. if os.path.exists("token.pickle"): @@ -81,11 +84,13 @@ def get_authenticated_service(): def get_monday_date(): + """Get Monday of current week. Weeks start on Monday.""" today = datetime.now() return today - timedelta(days=today.weekday()) -def create_playlist(youtube, title): +def make_playlist(youtube, title): + """Make a playlist with given title.""" response = ( youtube.playlists() .insert( @@ -93,7 +98,7 @@ def create_playlist(youtube, title): body={ "snippet": { "title": title, - "description": "Weekly playlist created by ParkerBot", + "description": "Weekly playlist generated by ParkerBot", }, "status": {"privacyStatus": "public"}, }, @@ -104,7 +109,8 @@ def create_playlist(youtube, title): return response["id"] -def get_or_create_playlist(youtube, monday_date): +def get_or_make_playlist(youtube, monday_date): + """Get ID of playlist for given Monday's week, make if doesn't exist.""" playlist_title = f"{PLAYLIST_TITLE} {monday_date.strftime('%Y-%m-%d')}" # Check if playlist exists in the database @@ -115,8 +121,8 @@ def get_or_create_playlist(youtube, monday_date): if row: return row[0] # Playlist already exists - # If not, create a new playlist on YouTube and save it in the database - playlist_id = create_playlist(youtube, playlist_title) + # If not, make a new playlist on YouTube and save it in the database + playlist_id = make_playlist(youtube, playlist_title) with conn: cursor.execute( "INSERT INTO playlists (title, playlist_id, creation_date) VALUES (?, ?, ?)", @@ -127,6 +133,7 @@ def get_or_create_playlist(youtube, monday_date): def add_video_to_playlist(youtube, playlist_id, video_id): + """Add video to playlist.""" youtube.playlistItems().insert( part="snippet", body={ @@ -138,32 +145,39 @@ def add_video_to_playlist(youtube, playlist_id, video_id): ).execute() -async def message_callback(room, event): +def is_music(youtube, video_id): + """Check whether a YouTube video is music.""" + video_details = youtube.videos().list(id=video_id, part="snippet").execute() + + # Check if the video category is Music (typically category ID 10) + return video_details["items"][0]["snippet"]["categoryId"] == "10" + + +async def message_callback(client, room, event): + """Event handler for received messages.""" youtube_link_pattern = r"(https?://(?:www\.|music\.)?youtube\.com/(?!playlist\?list=)watch\?v=[\w-]+|https?://youtu\.be/[\w-]+)" - if event.sender != MATRIX_USER: - sender = event.sender - message_body = event.body + sender = event.sender + if sender != MATRIX_USER: + body = event.body + timestamp = event.server_timestamp room_id = room.room_id monday_date = get_monday_date() youtube = get_authenticated_service() - playlist_id = get_or_create_playlist(youtube, monday_date) - youtube_links = re.findall(youtube_link_pattern, message_body) + playlist_id = get_or_make_playlist(youtube, monday_date) + youtube_links = re.findall(youtube_link_pattern, body) - if message_body == "!pow": + if body == "!pow": playlist_link = f"https://www.youtube.com/playlist?list={playlist_id}" reply_msg = f"{sender}, here's the playlist of the week: {playlist_link}" await client.room_send( room_id=room_id, message_type="m.room.message", - content={ - "msgtype": "m.text", - "body": reply_msg - } + content={"msgtype": "m.text", "body": reply_msg}, ) - if youtube_links: - timestamp = event.server_timestamp - for link in youtube_links: + for link in youtube_links: + video_id = link.split("v=")[-1] + if is_music(youtube, video_id): try: cursor.execute( "INSERT INTO messages (sender, message, timestamp) VALUES (?, ?, ?)", @@ -177,9 +191,6 @@ async def message_callback(room, event): else: raise e - for link in youtube_links: - video_id = link.split("v=")[-1] - # Check if the link is already added to any playlist cursor.execute("SELECT id FROM messages WHERE message = ?", (link,)) message_row = cursor.fetchone() @@ -204,15 +215,37 @@ async def message_callback(room, event): print(f"Added track to playlist: {link}") -async def main(): - create_tables() - global client +async def sync_callback(response): + # Save the sync token to a file or handle it as needed + with open("sync_token", "w") as f: + f.write(response.next_batch) + + +def load_sync_token(): + try: + with open("sync_token", "r") as file: + return file.read().strip() + except FileNotFoundError: + return None + + +async def get_client(): client = AsyncClient(MATRIX_SERVER, MATRIX_USER) - client.add_event_callback(message_callback, RoomMessageText) + client.add_event_callback( + lambda room, event: message_callback(client, room, event), RoomMessageText + ) + client.add_response_callback(sync_callback, SyncResponse) print(await client.login(MATRIX_PASSWORD)) await client.join(MATRIX_ROOM) - await client.sync_forever(timeout=10000) # milliseconds + return client +async def main(): + """Get DB and Matrix client ready, and start syncing.""" + define_tables() + client = await get_client() + sync_token = load_sync_token() + await client.sync_forever(30000, full_state=True, since=sync_token) + if __name__ == "__main__": asyncio.run(main())