TheHamkerCat hace 5 años
padre
commit
be116c3367

+ 8 - 2
wbb/modules/nsfw.py

@@ -38,7 +38,10 @@ async def detect_nsfw(_, message):
     async with aiohttp.ClientSession() as session:
         async with session.post("https://api.imgbb.com/1/upload", data=payload) as resp:
             data = await resp.json()
-        url = data['data']['url']
+        try:
+            url = data['data']['url']
+        except KeyError:
+            return
     os.remove(image)
     try:
         results = await arq.nsfw_scan(url)
@@ -104,7 +107,10 @@ async def nsfw_scan_command(_, message):
     async with aiohttp.ClientSession() as session:
         async with session.post("https://api.imgbb.com/1/upload", data=payload) as resp:
             data = await resp.json()
-        url = data['data']['url']
+        try:
+            url = data['data']['url']
+        except KeyError:
+            await m.edit("Failed to upload this to api server.")
     os.remove(image)
     try:
         results = await arq.nsfw_scan(url)

+ 1 - 1
wbb/modules/sudoers.py

@@ -64,7 +64,7 @@ async def bot_sys_stats():
 {USERBOT_USERNAME}@William
 ------------------
 UPTIME: {formatter.get_readable_time((bot_uptime))}
-BOT: {process.memory_info()[0] / 1024 ** 2} MB
+BOT: {round(process.memory_info()[0] / 1024 ** 2)} MB
 CPU: {cpu}%
 RAM: {mem}%
 DISK: {disk}%

+ 1 - 0
wbb/utils/aiodownloader/__init__.py

@@ -0,0 +1 @@
+from .downloader import Handler

+ 134 - 0
wbb/utils/aiodownloader/downloader.py

@@ -0,0 +1,134 @@
+import os
+import asyncio
+from typing import Optional, List
+
+import aiohttp
+import aiofiles
+
+
+class DownloadJob:
+    """
+    Download Job
+
+    :param file_url: url where the file is located
+    :param session: aiohttp session to be used on the job
+    :param file_name: name to be used on the file. Defaults to the last part of the url
+    :param save_path: dir where the file should be saved. Defaults to the current dir
+    """
+
+    def __init__(self,
+                 session: aiohttp.ClientSession,
+                 file_url: str,
+                 save_path: Optional[str] = None,
+                 chunk_size: Optional[int] = 1024):
+
+        self.file_url = file_url
+        self._session = session
+        self._chunk_size = chunk_size
+
+        self.file_name = file_url.split('/')[~0]
+        self.file_path = os.path.join(
+            save_path, self.file_name) if save_path else self.file_name
+
+        self.completed = False
+        self.progress = 0
+        self.size = 0
+
+    async def get_size(self) -> int:
+        """
+        Gets the content-length of the file from the file url.
+        :return: the files size in bytes
+        """
+        if not self.size:
+            async with self._session.get(self.file_url) as resp:
+                if 200 <= resp.status < 300:
+                    self.size = int(resp.headers['Content-Length'])
+
+                else:
+                    raise aiohttp.errors.HttpProcessingError(
+                        message=f'There was a problem processing {self.file_url}',
+                        code=resp.status)
+
+        return self.size
+
+    def _downloaded(self, chunk_size):
+        """
+        Method to be called when a chunk of the file is downloaded. It updates the
+        progress, adds the size to the _advance variable and checks if the download
+        is completed
+        """
+        self.progress += chunk_size
+
+    async def download(self):
+        """
+        Downloads the file from the given url to a file in the given path.
+        """
+
+        async with self._session.get(self.file_url) as resp:
+            # Checkning the response code
+            if 200 <= resp.status < 300:
+                # Saving the data to the file chunk by chunk.
+                async with aiofiles.open(self.file_path, 'wb') as file:
+
+                    # Downloading the file using the aiohttp.StreamReader
+                    async for data in resp.content.iter_chunked(self._chunk_size):
+                        await file.write(data)
+                        self._downloaded(self._chunk_size)
+
+                self.completed = True
+                return self
+
+            else:
+                raise aiohttp.errors.HttpProcessingError(
+                    message=f'There was a problem processing {self.file_url}', code=resp.status)
+
+
+class Handler:
+    """
+    Top level interface with the downloader. It creates the download jobs and handles them.
+
+    :param loop: asyncio loop. if not provided asyncio.get_event_loop() will be used
+    :param session: aiohttp session to be used on all downloads. If not provided it will be 
+    created 
+    :param chunk_size: chunk bytes sizes to get from the file source. Defaults to 1024 bytes 
+    """
+
+    def __init__(self,
+                 loop: Optional[asyncio.BaseEventLoop] = None,
+                 session: Optional[aiohttp.ClientSession] = None,
+                 chunk_size: Optional[int] = 1024):
+
+        self._loop = loop or asyncio.get_event_loop()
+        self._session = session or aiohttp.ClientSession(loop=self._loop)
+        self._chunk_size = chunk_size
+
+    def _job_factory(self,
+                     file_url: str,
+                     save_path: Optional[str] = None) -> DownloadJob:
+        """
+        Shortcut for creating a download job. It adds the session and the chunk size.
+        :param file_url: url where the file is located
+        :param save_path: save path for the download
+        :return: 
+        """
+
+        return DownloadJob(self._session, file_url, save_path, self._chunk_size)
+
+    async def download(self,
+                       *files_url: str,
+                       save_path: Optional[str] = None) -> List[DownloadJob]:
+        """
+        Downloads a bulk of files from the given list of urls to the given path.
+
+        :param files_url: list of urls where the files are located
+        :param save_path: path to be used for saving the files. Defaults to the current dir
+        """
+
+        jobs = [self._job_factory(url, save_path=save_path)
+                for url in files_url]
+
+        tasks = [asyncio.ensure_future(job.download()) for job in jobs]
+
+        await asyncio.gather(*tasks)
+
+        return [task.result() for task in tasks]

+ 8 - 0
wbb/utils/functions.py

@@ -23,9 +23,16 @@ SOFTWARE.
 """
 from PIL import Image, ImageDraw, ImageFont, ImageFilter
 from random import randint
+from wbb.utils import aiodownloader
 import speedtest
 import aiohttp
 
+"""
+Just import 'downloader' anywhere and do downloader.download() to
+download file from a given url
+"""
+downloader = aiodownloader.Handler()
+
 
 def generate_captcha():
     # Generate one letter
@@ -104,3 +111,4 @@ async def get_http_status_code(url: str) -> int:
     async with aiohttp.ClientSession() as session:
         async with session.head(url) as resp:
             return resp.status
+

+ 21 - 16
wbb/utils/inlinefuncs.py

@@ -122,19 +122,25 @@ async def alive_function(answers):
 async def translate_func(answers, lang, tex):
     i = Translator().translate(tex, dest=lang)
     msg = f"""
-__**Translated to {lang}**__
+__**Translated from {i.src} to {lang}**__
 
 **INPUT:**
 {tex}
 
 **OUTPUT:**
 {i.text}"""
-    answers.append(
-        InlineQueryResultArticle(
-            title=f'Translated to {lang}',
-            description=i.text,
-            input_message_content=InputTextMessageContent(msg)
-        )
+    answers.extend(
+        [
+            InlineQueryResultArticle(
+                title=f'Translated from {i.src} to {lang}.',
+                description=i.text,
+                input_message_content=InputTextMessageContent(msg)
+            ),
+            InlineQueryResultArticle(
+                title=i.text,
+                input_message_content=InputTextMessageContent(i.text)
+            )
+        ]
     )
     return answers
 
@@ -300,7 +306,7 @@ async def shortify(url):
     async with aiohttp.ClientSession() as session:
         async with session.post("https://api-ssl.bitly.com/v4/shorten", headers=header, data=payload) as resp:
             data = await resp.json()
-    msg = f"**Original Url:** {url}\n**Shortened Url:** {data['link']}"
+    msg = data['link']
     a = []
     b = InlineQueryResultArticle(
         title="Link Shortened!",
@@ -797,7 +803,7 @@ async def ping_func(answers):
     return answers
 
 
-async def nsfw_scan_func(answers, url:str):
+async def nsfw_scan_func(answers, url: str):
     t1 = time()
     data = (await arq.nsfw_scan(url)).data
     t2 = time()
@@ -812,11 +818,10 @@ async def nsfw_scan_func(answers, url:str):
 **Sexy:** `{data.sexy} %`
     """
     answers.append(
-            InlineQueryResultArticle(
-                title="Scanned",
-                description=f"Took {tt} Seconds.",
-                input_message_content=InputTextMessageContent(content)
-                )
-            )
+        InlineQueryResultArticle(
+            title="Scanned",
+            description=f"Took {tt} Seconds.",
+            input_message_content=InputTextMessageContent(content)
+        )
+    )
     return answers
-