Save sync token, check if videos are music before saving.

Signed-off-by: Abdulkadir Furkan Şanlı <me@abdulocra.cy>

.gitignore

@@ -5,3 +5,4 @@ venv
 *.pickle
 *.json
 *.swp
+sync_token

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())