tea/tg_gentoo_bot
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

third version

Browse Source 
* added ratelimiting and task grouping
* bumped python to 3.11
This commit is contained in:
Aleksey Chubukov 2023-02-07 09:06:58 +03:00
parent ec8008c29e
commit d9fb692eeb
2 changed files with 74 additions and 65 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
Pipfile
View File

@ -10,7 +10,7 @@ telethon = "*"
[dev-packages]
[requires]
python_version = "3.9"
python_version = "3.11"
[scripts]
bot = "python3 ./bot.py"

89
bot.py
View File

@ -6,10 +6,9 @@ from telethon import functions, types, errors
from pathlib import Path
import asyncio
logging.basicConfig(filename=str('bot.log'),level=logging.DEBUG)
logging.basicConfig(filename=str('bot.log'),level=logging.INFO)
log = logging.getLogger('gentoobot')
semaphore = asyncio.Semaphore(10)
class StickerPackDownloader:
_log = logging.getLogger('gentoobot.packdownloader')
@ -23,18 +22,21 @@ class StickerPackDownloader:
self.sticker_downloads.mkdir(exist_ok=True)
async def download(self, inputpack):
if inputpack.id not in self.seen:
if inputpack.id in self.seen:
return
self.seen.add(inputpack.id)
try:
req = functions.messages.GetStickerSetRequest(stickerset=inputpack,hash=0)
await for_online()
stickers = await self.client(req)
self._log.info("downloadding stickerset %s %s", inputpack.id, stickers.set.title)
sticker_dir = self.sticker_downloads/str(inputpack.id)
sticker_dir.mkdir(exist_ok=True)
with open(str(sticker_dir/'meta.txt'), 'w') as meta:
meta.write(stickers.to_json())
self._log.info("preparing to download whole stickerset %s", inputpack.id)
self._log.debug("preparing to download whole stickerset %s", inputpack.id)
for doc in stickers.documents:
await self.downloader.download(doc)
self.seen.add(inputpack.id)
except errors.rpcerrorlist.StickersetInvalidError:
self._log.warning("sadly, stickerset %s no longer exists", inputpack.id)
@ -52,7 +54,7 @@ class StickerDownloader:
self.sticker_downloads = self.downloads_root/'stickers'
self.sticker_downloads.mkdir(exist_ok=True)
self.seen_stickers = set()
self.seen_packs = set()
self.pending_stickers = set()
self.client = client
@classmethod
@ -71,21 +73,30 @@ class StickerDownloader:
meta['alt'] = a.alt
return meta
async def download(self, document):
if document.id not in self.seen_stickers:
async def download_task(self, document):
if document.id in self.seen_stickers:
return
self.seen_stickers.add(document.id)
try:
meta = self.get_sticker_metadata(document)
to_dir = self.sticker_downloads / str(meta.get('pack', 'inline'))
to_dir.mkdir(exist_ok=True)
to_file = to_dir / (str(document.id)+'.'+self.mimes.get(document.mime_type, ''))
if not to_file.exists() or to_file.stat().st_size != document.size:
self._log.info("downloading %s to %s", document.id, str(to_file))
self._log.debug("downloading %s to %s", document.id, str(to_file))
await for_online()
async with semaphore:
await self.client.download_file(document, to_file)
self.seen_stickers.add(document.id)
self._log.info("downloaded %s to %s", document.id, str(to_file))
else:
self._log.info("document %s already downloaded at %s", document.id, str(to_file))
self._log.debug("document %s already downloaded at %s", document.id, str(to_file))
except Exception as e:
self._log.critical("oops: ", exc_info=e)
self.seen_stickers.remove(document.id)
async def download(self, document):
asyncio.create_task(self.download_task(document))
@ -103,36 +114,21 @@ def make_client(nth):
log.debug("client id %s hash %s", app_id, app_hash)
return TelegramClient('session'+str(nth)+str(app_id), app_id, app_hash)
def main():
cfg = load_config('config.yaml')
app_id = cfg['apps'][0]['id']
app_hash = cfg['apps'][0]['hash']
log.debug("client id %s hash %s", app_id, app_hash)
client = TelegramClient('session0'+str(app_id), app_id, app_hash)
client.start()
log.debug(client.get_me().stringify())
dl = StickerDownloader(client, './download')
try:
for found in iter_sources(client, cfg['sources']):
if isinstance(found, types.Document):
dl.download(found)
elif isinstance(found, types.InputStickerSetID):
dl.download_pack(found)
except KeyboardInterrupt:
log.warn('exiting: interrupted by keyboard')
client.disconnect()
class DialogProcessor:
_log = logging.getLogger('gentoobot.messages')
def __init__(self, client, dialog, downloader, packdownloader):
def __init__(self, client, dialog, task_group, destdir):
self.client = client
self.downloader = downloader
self.packdownloader = packdownloader
self.dialog = dialog
self.task_group = task_group
self.downloader = StickerDownloader(self.client, destdir)
self.packdownloader = StickerPackDownloader(self.downloader)
async def process(self):
async for message in self.client.iter_messages(self.dialog):
asyncio.create_task(self.check_message(message))
await for_online()
async with semaphore:
self._log.debug("starting task for message")
self.task_group.create_task(self.check_message(message))
async def check_message(self, message):
try:
@ -145,9 +141,11 @@ class DialogProcessor:
for a in message.media.document.attributes:
if hasattr(a, 'stickerset'):
self._log.debug("document %s is a sticker", message.media.document.id)
await for_online()
await self.downloader.download(message.media.document)
if isinstance(a.stickerset, types.InputStickerSetID):
self._log.debug("document %s belongs to stickerset", message.media.document.id)
await for_online()
await self.packdownloader.download(a.stickerset)
except Exception as e:
self._log.error("somethin wrong happened during checking message: %s", message.stringify(), exc_info=e)
@ -157,6 +155,7 @@ class ArchiveWalker:
def __init__(self, client, archives):
self.client = client
self.archives = archives
self.task_group = asyncio.TaskGroup()
async def fetch_dialogs(self):
async for dialog in self.client.iter_dialogs():
@ -176,17 +175,27 @@ class ArchiveWalker:
yield dialog, src
async def walk(self):
async with self.task_group:
async for dialog, options in self.fetch_dialogs():
downloader = StickerDownloader(self.client, Path(options.get('destdir', 'downloads')))
packdownloader = StickerPackDownloader(downloader)
processor = DialogProcessor(self.client, dialog, downloader, packdownloader)
asyncio.create_task(processor.process())
processor = DialogProcessor(self.client, dialog, self.task_group, Path(options.get('destdir', 'downloads')))
self._log.debug('starting dialog processor')
self.task_group.create_task(processor.process())
self._log.debug('walker finished')
async def for_online():
while not client.is_connected():
log.warn('disconnected')
await asyncio.sleep(10)
async def main():
try:
log.debug((await client.get_me()).stringify())
walker = ArchiveWalker(client, cfg['sources'])
await walker.walk()
await client.run_until_disconnected()
except KeyboardInterrupt:
return