Add estimator provider to make more relevant calls

bot/handlers/tiktok/on_tiktok.py

@@ -1,3 +1,5 @@
+import time
+
 from aiogram import Router, types
 from rich import print
 from rich.traceback import Traceback
@@ -5,7 +7,7 @@ from rich.traceback import Traceback
 from bot.filters import TikTokUrlFilter
 from bot.modules.custom_sender import custom_sender
 from bot.modules.cycle.cycle import LinkCycle
-from bot.modules.cycle.estimator import process_situation
+from bot.modules.cycle.processor import process_situation
 from bot.utils.config import config
 
 router = Router()
@@ -46,6 +48,7 @@ async def on_tiktok(message: types.Message):
                         await message.reply_video(
                             video=types.URLInputFile(url=cycle.video.url)
                         )
+
                 if cycle.photos:
                     for chunk in cycle.photos.urls_chunked:
                         if len(chunk) > 1:

bot/modules/cycle/cycle.py

@@ -40,7 +40,7 @@ class LinkCycle:
 
     job: Jobs = field(default_factory=Jobs)
     needs: PropagatePhotosNeeded | PropagateVideoNeeded | PropagateEverythingNeeded | PropagateAudioNeeded = (
-        None
+        field(default_factory=PropagateEverythingNeeded)
     )
 
     video: BaseVideo = None

bot/modules/cycle/processor.py

@@ -49,11 +49,11 @@ def __process_estimated(cycle: LinkCycle):
 
 async def process_situation(cycle: LinkCycle) -> LinkCycle:
     try:
-        if cycle.job.estimate and not cycle.needs:
+        if cycle.job.estimate and not cycle.needs.needs:
             cycle.estimated = await cycle.providers.get_next().from_url(cycle.url)
             return __process_estimated(cycle)
 
-        elif cycle.job.estimate and cycle.needs:
+        elif cycle.job.estimate and cycle.needs.needs:
             provider = cycle.providers.get_for_type(cycle.needs.needs)
             cycle.estimated = await (
                 provider.from_preferred(getattr(cycle, provider.wants))

bot/modules/providers/__init__.py

@@ -3,6 +3,7 @@ from dataclasses import dataclass
 from .common.providers import Provider
 from .contentstudio import contentstudio
 from .tiktokapi import tiktokapi
+from .estimator import estimator
 
 
 class NoProvidersLeft(Exception):
@@ -11,7 +12,7 @@ class NoProvidersLeft(Exception):
 
 @dataclass
 class Providers:
-    __providers = [contentstudio, tiktokapi]
+    __providers = [estimator, contentstudio, tiktokapi]
     __current = -1
 
     def get_next(self) -> Provider:
@@ -22,9 +23,8 @@ class Providers:
             raise NoProvidersLeft()
 
     def get_for_type(self, supports):
-        if supports is None:
-            return self.get_next()
         try:
+            self.__current += 1
             current = self.__providers[self.__current]
         except IndexError:
             raise NoProvidersLeft()

bot/modules/providers/common/propagations.py

@@ -4,7 +4,7 @@ from typing import Optional
 
 @dataclass
 class PropagatePhotosNeeded:
-    audio: str
+    audio: str = None
     id: Optional[str] = None
     needs = "photos"
 

bot/modules/providers/estimator/__init__.py

@@ -0,0 +1,6 @@
+from .estimator import Estimator
+
+estimator = Estimator()
+
+
+__all__ = ["estimator"]

bot/modules/providers/estimator/content.py

@@ -0,0 +1,46 @@
+from typing import Optional
+
+from attrs import define
+
+from ..common.propagations import PropagateVideoNeeded, PropagatePhotosNeeded, PropagateEverythingNeeded
+
+import re
+
+REGEX_VIDEO_ID = (
+    r"(?:https:\/\/(?:www\.)*tiktok\.com\/@[^?\/]+\/video\/)(?:([0-9]+)?(?:\?.+)?$|$)"
+)
+REGEX_PHOTO_ID = (
+    r"(?:https:\/\/(?:www\.)*tiktok\.com\/@[^?\/]+\/photo\/)(?:([0-9]+)?(?:\?.+)?$|$)"
+)
+REGEX_TIKTOK_URL = (
+    r"(https:\/\/(?:www\.)*tiktok\.com\/@[^?\/]+)"
+    r"(?:(\/(?:video|photo)\/[0-9]+)?(\?.+)?$|$)"
+)
+
+
+@define
+class Content:
+    @staticmethod
+    async def from_id(
+        video_id: str
+    ) -> Optional[PropagateEverythingNeeded]:
+        return PropagateEverythingNeeded(video_id)
+
+    @staticmethod
+    async def from_url(
+        url: str
+    ) -> Optional[PropagateVideoNeeded | PropagatePhotosNeeded | PropagateEverythingNeeded]:
+        if g := re.search(REGEX_VIDEO_ID, url, re.IGNORECASE):
+            return PropagateVideoNeeded(
+                id=g.groups()[0]
+            )
+        elif g := re.search(REGEX_PHOTO_ID, url, re.IGNORECASE):
+            return PropagatePhotosNeeded(
+                id=g.groups()[0]
+            )
+        elif g := re.search(REGEX_TIKTOK_URL, url, re.IGNORECASE):
+            return PropagateEverythingNeeded(
+                id=g.groups()[1].split("/")[-1]
+            )
+        else:
+            return PropagateEverythingNeeded()

bot/modules/providers/estimator/estimator.py

@@ -0,0 +1,53 @@
+from async_lru import alru_cache
+
+from ..common.propagations import (
+    PropagateAudioNeeded,
+    PropagateEverythingNeeded,
+    PropagatePhotosNeeded,
+    PropagateVideoNeeded,
+)
+from ..common.providers import Provider
+from .content import Content
+
+
+class Estimator(Provider):
+    supports_video = True
+    supports_audio = True
+    supports_photos = True
+    wants = "url"
+
+    def __init__(self):
+        self.content = Content()
+
+    @alru_cache()
+    async def from_url(
+        self, url: str
+    ) -> (
+            PropagateEverythingNeeded
+            | PropagateAudioNeeded
+            | PropagateVideoNeeded
+            | PropagatePhotosNeeded
+    ):
+        return await self.content.from_url(url)
+
+    @alru_cache()
+    async def from_id(
+        self, url: str
+    ) -> (
+            PropagateEverythingNeeded
+            | PropagateAudioNeeded
+            | PropagateVideoNeeded
+            | PropagatePhotosNeeded
+    ):
+        return await self.content.from_id(url)
+
+    @alru_cache()
+    async def from_preferred(
+        self, preferred: str
+    ) -> (
+            PropagateEverythingNeeded
+            | PropagateAudioNeeded
+            | PropagateVideoNeeded
+            | PropagatePhotosNeeded
+    ):
+        return await self.from_url(preferred)