From 6c64fcd75379428c3e113699be01ee7f264505b3 Mon Sep 17 00:00:00 2001 From: hhh Date: Wed, 13 Mar 2024 22:57:51 +0200 Subject: [PATCH] Add estimator provider to make more relevant calls --- bot/handlers/tiktok/on_tiktok.py | 5 +- bot/modules/cycle/cycle.py | 2 +- .../cycle/{estimator.py => processor.py} | 4 +- bot/modules/providers/__init__.py | 6 +-- bot/modules/providers/common/propagations.py | 2 +- bot/modules/providers/estimator/__init__.py | 6 +++ bot/modules/providers/estimator/content.py | 46 ++++++++++++++++ bot/modules/providers/estimator/estimator.py | 53 +++++++++++++++++++ 8 files changed, 116 insertions(+), 8 deletions(-) rename bot/modules/cycle/{estimator.py => processor.py} (95%) create mode 100644 bot/modules/providers/estimator/__init__.py create mode 100644 bot/modules/providers/estimator/content.py create mode 100644 bot/modules/providers/estimator/estimator.py diff --git a/bot/handlers/tiktok/on_tiktok.py b/bot/handlers/tiktok/on_tiktok.py index 0736add..c41fc62 100644 --- a/bot/handlers/tiktok/on_tiktok.py +++ b/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: diff --git a/bot/modules/cycle/cycle.py b/bot/modules/cycle/cycle.py index 740da81..13f06a4 100644 --- a/bot/modules/cycle/cycle.py +++ b/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 diff --git a/bot/modules/cycle/estimator.py b/bot/modules/cycle/processor.py similarity index 95% rename from bot/modules/cycle/estimator.py rename to bot/modules/cycle/processor.py index 44500ba..bd8ca1f 100644 --- a/bot/modules/cycle/estimator.py +++ b/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)) diff --git a/bot/modules/providers/__init__.py b/bot/modules/providers/__init__.py index e0bfc2f..a688e39 100644 --- a/bot/modules/providers/__init__.py +++ b/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() diff --git a/bot/modules/providers/common/propagations.py b/bot/modules/providers/common/propagations.py index c646b98..da95af8 100644 --- a/bot/modules/providers/common/propagations.py +++ b/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" diff --git a/bot/modules/providers/estimator/__init__.py b/bot/modules/providers/estimator/__init__.py new file mode 100644 index 0000000..358a398 --- /dev/null +++ b/bot/modules/providers/estimator/__init__.py @@ -0,0 +1,6 @@ +from .estimator import Estimator + +estimator = Estimator() + + +__all__ = ["estimator"] diff --git a/bot/modules/providers/estimator/content.py b/bot/modules/providers/estimator/content.py new file mode 100644 index 0000000..291231a --- /dev/null +++ b/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() diff --git a/bot/modules/providers/estimator/estimator.py b/bot/modules/providers/estimator/estimator.py new file mode 100644 index 0000000..b446332 --- /dev/null +++ b/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)