Add downloading

This commit is contained in:
2022-01-09 20:12:06 +03:00
parent f062b41ce8
commit 42ebed5503
5 changed files with 135 additions and 31 deletions

View File

@@ -0,0 +1,61 @@
from io import BytesIO
from typing import Optional
from telegram_files_storage import AiogramFilesStorage, TelethonFilesStorage
from app.services.storages import StoragesContainer
class FileDownloader:
_aiogram_storage_index = 0
_telethon_storage_index = 0
@classmethod
@property
def AIOGRAM_STORAGES(cls) -> list[AiogramFilesStorage]:
return StoragesContainer.AIOGRAM_STORAGES
@classmethod
@property
def TELETHON_STORAGES(cls) -> list[TelethonFilesStorage]:
return StoragesContainer.TELETHON_STORAGES
@classmethod
def get_aiogram_storage(cls) -> AiogramFilesStorage:
if not cls.AIOGRAM_STORAGES:
raise ValueError("Aiogram storage not exist!")
cls._aiogram_storage_index = (cls._aiogram_storage_index + 1) % len(
cls.AIOGRAM_STORAGES
)
return cls.AIOGRAM_STORAGES[cls._aiogram_storage_index]
@classmethod
def get_telethon_storage(cls) -> TelethonFilesStorage:
if not cls.TELETHON_STORAGES:
raise ValueError("Telethon storage not exists!")
cls._telethon_storage_index = (cls._telethon_storage_index + 1) % len(
cls.TELETHON_STORAGES
)
return cls.TELETHON_STORAGES[cls._telethon_storage_index]
@classmethod
async def download_by_file_id(cls, file_id: str) -> Optional[BytesIO]:
if not cls.AIOGRAM_STORAGES:
return None
storage = cls.get_aiogram_storage()
return await storage.download(file_id)
@classmethod
async def download_by_message_id(cls, message_id: int) -> Optional[BytesIO]:
if not cls.TELETHON_STORAGES:
return None
storage = cls.get_telethon_storage()
return await storage.download(message_id)

View File

@@ -6,16 +6,23 @@ from fastapi import UploadFile
from telegram_files_storage import AiogramFilesStorage, TelethonFilesStorage
from app.models import UploadedFile, UploadBackends
from core.config import env_config
from app.services.storages import StoragesContainer
class FileUploader:
AIOGRAM_STORAGES: list[AiogramFilesStorage] = []
TELETHON_STORAGES: list[TelethonFilesStorage] = []
_aiogram_storage_index = 0
_telethon_storage_index = 0
@classmethod
@property
def AIOGRAM_STORAGES(cls) -> list[AiogramFilesStorage]:
return StoragesContainer.AIOGRAM_STORAGES
@classmethod
@property
def TELETHON_STORAGES(cls) -> list[TelethonFilesStorage]:
return StoragesContainer.TELETHON_STORAGES
def __init__(self, file: UploadFile, caption: Optional[str] = None) -> None:
self.file = file
self.caption = caption
@@ -81,28 +88,6 @@ class FileUploader:
data=self.upload_data,
)
@classmethod
async def prepare(cls):
if env_config.BOT_TOKENS:
cls.AIOGRAM_STORAGES: list[AiogramFilesStorage] = [
AiogramFilesStorage(env_config.TELEGRAM_CHAT_ID, token)
for token in env_config.BOT_TOKENS
]
if env_config.TELETHON_APP_CONFIG and env_config.TELETHON_SESSIONS:
cls.TELETHON_STORAGES: list[TelethonFilesStorage] = [
TelethonFilesStorage(
env_config.TELEGRAM_CHAT_ID,
env_config.TELETHON_APP_CONFIG.APP_ID,
env_config.TELETHON_APP_CONFIG.API_HASH,
session,
)
for session in env_config.TELETHON_SESSIONS
]
for storage in [*cls.AIOGRAM_STORAGES, *cls.TELETHON_STORAGES]:
await storage.prepare()
@classmethod
def get_aiogram_storage(cls) -> AiogramFilesStorage:
if not cls.AIOGRAM_STORAGES:

View File

@@ -0,0 +1,30 @@
from telegram_files_storage import AiogramFilesStorage, TelethonFilesStorage
from core.config import env_config
class StoragesContainer:
AIOGRAM_STORAGES: list[AiogramFilesStorage] = []
TELETHON_STORAGES: list[TelethonFilesStorage] = []
@classmethod
async def prepare(cls):
if env_config.BOT_TOKENS:
cls.AIOGRAM_STORAGES: list[AiogramFilesStorage] = [
AiogramFilesStorage(env_config.TELEGRAM_CHAT_ID, token)
for token in env_config.BOT_TOKENS
]
if env_config.TELETHON_APP_CONFIG and env_config.TELETHON_SESSIONS:
cls.TELETHON_STORAGES: list[TelethonFilesStorage] = [
TelethonFilesStorage(
env_config.TELEGRAM_CHAT_ID,
env_config.TELETHON_APP_CONFIG.APP_ID,
env_config.TELETHON_APP_CONFIG.API_HASH,
session,
)
for session in env_config.TELETHON_SESSIONS
]
for storage in [*cls.AIOGRAM_STORAGES, *cls.TELETHON_STORAGES]:
await storage.prepare()