fresh start

Twitter/downloader.py

@@ -0,0 +1,147 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+import asyncio
+from datetime import datetime
+import gc
+import traceback
+import tracemalloc
+
+from Classifier.classifyHelper import classify_all
+from Database.x_classes import DownloadMode
+from Discord import discordHelper
+from Twitter import tweetHelper
+from tweety.types import Tweet
+from exceptions import ACCOUNT_DEAD, ACCOUNT_SKIP, DOWNLOAD_FAIL, NO_CHANNEL, OTHER_ERROR
+from Database.x_classes import ActionTaken, DownloadMode, ErrorID, HavoxLabel, PostRating, x_posts, x_posts_images, x_accounts
+from Database.db_schema import x_accounts as schema_x_accounts
+if TYPE_CHECKING:
+    from runtimeBotData import RuntimeBotData
+
+async def download_loop(botData: RuntimeBotData):
+    guild = botData.client.guilds[0]
+    try:
+        results = {}
+        botData.new_accounts = []
+        
+        for artist in botData.db.x_get_all_accounts():
+            if artist.is_deleted: continue
+
+            #sleep to avoid rate limits
+            await asyncio.sleep(5)
+
+            print("Artist:", artist.name)
+
+            #wait for ALL new posts to be found
+            try:
+                match artist.download_mode:
+                    case DownloadMode.NO_DOWNLOAD:
+                        continue
+                    case DownloadMode.DOWNLOAD:
+                        await discordHelper.ensure_has_channel_or_thread(artist, guild, botData.db)
+                        new_posts = await tweetHelper.UpdateMediaPosts(artist, botData)
+                    case DownloadMode.DOWNLOAD_ALL:
+                        await discordHelper.ensure_has_channel_or_thread(artist, guild, botData.db)
+                        new_posts = await tweetHelper.DownloadAllMediaPosts(artist, botData)
+                    case _:
+                        continue
+            except ACCOUNT_DEAD:
+                botData.db.x_update_account_properties(artist.id, [(schema_x_accounts.is_deleted, True)])
+                continue
+            except ACCOUNT_SKIP:
+                continue
+
+            if len(new_posts) == 0: continue
+
+            new_posts_count = len([post for post in new_posts if len(post.media) > 0])
+            if new_posts_count > 20:
+                #skips posting to discord if there are too many posts
+                botData.new_accounts.append(artist.name)
+
+            new_posts.sort(key= lambda x: x.date)
+
+            for tweet in new_posts: #posts should arrive here in chronological order
+                await download_post(artist, tweet, botData)
+                gc.collect()
+                print(tracemalloc.get_traced_memory())
+
+            results[artist.name] = new_posts_count
+            if artist.download_mode == DownloadMode.DOWNLOAD_ALL:
+                botData.db.x_update_account_properties(artist.id, [(schema_x_accounts.download_mode, DownloadMode.DOWNLOAD)])
+        await discordHelper.post_result(results, guild, botData.new_accounts)
+    except Exception as ex:
+        print(ex)
+        await discordHelper.send_error(traceback.format_exc()[0:256], botData)
+ 
+async def download_post(artist: x_accounts, tweet: Tweet, botData: RuntimeBotData):
+    x_post = x_posts(id = tweet.id, account_id = tweet.author.id, date = tweet.date, text = tweet.text)
+
+    if len(tweet.media) == 0:
+        x_post.error_id = ErrorID.NO_ART
+        botData.db.x_insert_post(x_post, commit=True)
+        return
+    
+    print("New media post:", str(tweet.url))
+    media = await tweetHelper.GetTweetMediaUrls(tweet)
+    image_containers = [x_posts_images(tweet.id, idx, file = url) for idx, url in enumerate(media)]
+
+    try:
+        downloaded_media = await tweetHelper.DownloadMedia(tweet.id, tweet.author.id, tweet.author.username, media, botData.session)
+    except DOWNLOAD_FAIL as e:
+        x_post.error_id = e.code
+        
+        botData.db.x_insert_post(x_post, commit=False)
+        for image in image_containers:
+            image.error_id = ErrorID.DOWNLOAD_FAIL
+            botData.db.x_insert_image(image, commit=False)
+        botData.db.conn.commit()
+        return
+
+    def get_rating_value(rating):
+        return 4 if rating == PostRating.Explicit else 3 if rating == PostRating.Questionable else 2 if rating == PostRating.Sensitive else 1 if rating == PostRating.General else 0
+
+    vox_labels = []
+    final_filtered_tags = {}
+    duplicates = []
+    for idx, attachment in enumerate(downloaded_media):
+        container = image_containers[idx]
+        container.saved_file = attachment.file_name
+        container.vox_label, container.rating, container.tags, filtered_tags, container.phash, container.dhash, container.error_id = await classify_all(attachment.file_bytes, botData.classifier, botData.vox)
+        
+        if container.vox_label not in vox_labels:
+            vox_labels.append(container.vox_label)
+
+        if container.phash != None:
+            duplicate = botData.db.x_search_duplicate(user_id=x_post.account_id, max_id = x_post.id, phash=container.phash)
+            if duplicate != None:
+                container.duplicate_id = duplicate.post_id
+                container.duplicate_index = duplicate.index
+                if duplicate.post_id not in duplicates:
+                    duplicates.append(duplicate.post_id)
+
+        x_post.tags         = list(set(x_post.tags + container.tags))
+        x_post.rating       = container.rating if get_rating_value(container.rating) > get_rating_value(x_post.rating) else x_post.rating
+        final_filtered_tags = final_filtered_tags | filtered_tags
+
+    is_filtered = len(vox_labels) == 1 and vox_labels[0] == HavoxLabel.Rejected
+
+    try:
+        discord_post = await discordHelper.send_x_post(tweet, artist, botData.client.guilds[0], botData.new_accounts, downloaded_media, is_filtered, rating=x_post.rating, tags=final_filtered_tags, vox_labels = vox_labels, duplicate_posts=duplicates, xView=botData.xView, yView=botData.yView)
+    except (NO_CHANNEL, OTHER_ERROR) as e:
+        x_post.error_id = e.code
+        x_post.discord_post_id = 0
+    else:
+        x_post.discord_post_id = discord_post.id
+        
+    x_post.action_taken = ActionTaken.Rejected if is_filtered else ActionTaken.Null
+
+    try:
+        if not botData.db.x_insert_post(x_post, commit = False):
+            raise Exception("Transaction error")
+        
+        for image in image_containers:
+            botData.db.x_insert_image(image, False)
+    except Exception as ex:
+        botData.db.conn.rollback()
+        raise ex
+    else:
+        botData.db.conn.commit()