gd_count.py

#!/usr/bin/env python3
from pyrogram.handlers import MessageHandler
from pyrogram.filters import command

from bot import bot
from bot.helper.mirror_utils.upload_utils.gdriveTools import GoogleDriveHelper
from bot.helper.telegram_helper.message_utils import deleteMessage, sendMessage
from bot.helper.telegram_helper.filters import CustomFilters
from bot.helper.telegram_helper.bot_commands import BotCommands
from bot.helper.ext_utils.bot_utils import is_gdrive_link, sync_to_async, new_task, get_readable_file_size
from bot.helper.themes import BotTheme


@new_task
async def countNode(_, message):
    args = message.text.split()
    if username := message.from_user.username:
        tag = f"@{username}"
    else:
        tag = message.from_user.mention

    link = args[1] if len(args) > 1 else ''
    if len(link) == 0 and (reply_to := message.reply_to_message):
        link = reply_to.text.split(maxsplit=1)[0].strip()

    if is_gdrive_link(link):
        msg = await sendMessage(message, BotTheme('COUNT_MSG', LINK=link))
        gd = GoogleDriveHelper()
        name, mime_type, size, files, folders = await sync_to_async(gd.count, link)
        if mime_type is None:
            await sendMessage(message, name)
            return
        await deleteMessage(msg)
        msg = BotTheme('COUNT_NAME', COUNT_NAME=name)
        msg += BotTheme('COUNT_SIZE', COUNT_SIZE=get_readable_file_size(size))
        msg += BotTheme('COUNT_TYPE', COUNT_TYPE=mime_type)
        if mime_type == 'Folder':
            msg += BotTheme('COUNT_SUB', COUNT_SUB=folders)
            msg += BotTheme('COUNT_FILE', COUNT_FILE=files)
        msg += BotTheme('COUNT_CC', COUNT_CC=tag)
    else:
        msg = 'Send Gdrive link along with command or by replying to the link by command'
    await sendMessage(message, msg, photo='IMAGES')


bot.add_handler(MessageHandler(countNode, filters=command(
    BotCommands.CountCommand) & CustomFilters.authorized & ~CustomFilters.blacklisted))