refactor + fixes add watermark on images #3

telewater/bot.py

@@ -6,9 +6,9 @@
 
 from telethon import events
 
-from telewater.const import COMMANDS, HELP, Config, config
-from telewater.utils import download_image, get_args
-from telewater.watermark import watermark_video
+from telewater import conf
+from telewater.utils import cleanup, download_image, get_args, stamp
+from telewater.watermark import watermark_image, watermark_video
 
 
 async def start(event):
@@ -18,7 +18,7 @@ async def start(event):
 
 async def bot_help(event):
     try:
-        await event.respond(HELP)
+        await event.respond(conf.HELP)
     finally:
         raise events.StopPropagation
 
@@ -35,17 +35,19 @@ async def set_config(event):
 
         key, value = [item.strip() for item in splitted]
 
-        config_dict = config.dict()
+        config_dict = conf.config.dict()
         if not key in config_dict.keys():
-            raise ValueError(f"The key {key} is not a valid key in configuration.")
+            raise ValueError(
+                f"The key {key} is not a valid key in configuration.")
 
         config_dict[key] = value
         print(config_dict)
 
-        config = Config(**config_dict)
+        conf.config = conf.Config(**config_dict)
 
-        print(config)
+        print(conf.config)
         if key == "watermark":
+            cleanup("image.png")
             download_image(url=value)
         await event.respond(f"The value of {key} was set to {value}")
 
@@ -61,10 +63,9 @@ async def set_config(event):
 
 async def get_config(event):
     """usage /get KEY"""
-    global config
     try:
         key = get_args(event.message.text)
-        config_dict = config.dict()
+        config_dict = conf.config.dict()
         await event.respond(f"{config_dict.get(key)}")
     except ValueError as err:
         print(err)
@@ -76,21 +77,18 @@ async def get_config(event):
 
 
 async def watermarker(event):
-    # TODO: reject large files (above certain file limit)
-    # TODO: also watermark photos
     global config
     if event.gif or event.video:
-
-        mp4_file = await event.download_media("")
-        # TODO: suffix the downloaded media with time-stamp and user id
-
-        outf = watermark_video(mp4_file)
-        print(outf)
-        await event.client.send_file(event.sender_id, outf)
-        os.remove(mp4_file)
-        os.remove(outf)
+        watermark = watermark_video
     elif event.photo:
-        await event.respond("Photos are currently not supported")
+        watermark = watermark_image
+    else:
+        return
+
+    org_file = stamp(await event.download_media(""), user=str(event.sender_id))
+    out_file = watermark(org_file)
+    await event.client.send_file(event.sender_id, out_file)
+    cleanup(org_file, out_file)
 
 
 ALL_EVENTS = {

telewater/main.py

@@ -6,25 +6,24 @@
 
 from telethon.sync import TelegramClient, functions, types
 
+from telewater import conf
 from telewater.bot import ALL_EVENTS
-from telewater.const import COMMANDS, config
 from telewater.utils import download_image
 
 
 def start_bot(API_ID: int, API_HASH: str, name: str, token: str):
     os.makedirs(name, exist_ok=True)
     os.chdir(name)
 
-    if config.watermark:
-        download_image(url=config.watermark)
+    download_image(url=conf.config.watermark)
 
     client = TelegramClient(name, API_ID, API_HASH).start(bot_token=token)
 
     client(
         functions.bots.SetBotCommandsRequest(
             commands=[
                 types.BotCommand(command=key, description=value)
-                for key, value in COMMANDS.items()
+                for key, value in conf.COMMANDS.items()
             ]
         )
     )

telewater/utils.py

@@ -2,12 +2,19 @@
 """
 
 
+import logging
+import os
+import re
 import shutil
+from datetime import datetime
 
 import requests
 
 
 def download_image(url: str, filename: str = "image.png") -> bool:
+    if filename in os.listdir():
+        print("Image exists")
+        return True
     try:
         print("Downloading image ... ")
         response = requests.get(url, stream=True)
@@ -34,3 +41,26 @@ def get_args(text: str):
     args = args.strip()
     print(args)
     return args
+
+
+def cleanup(*files):
+    for file in files:
+        try:
+            os.remove(file)
+        except FileNotFoundError:
+            logging.info(f"File {file} does not exist.")
+
+
+def stamp(file: str, user: str):
+
+    now = str(datetime.now())
+    outf = safe_name(f"{user} {now} {file}")
+    try:
+        os.rename(file, outf)
+        return outf
+    except Exception as err:
+        logging.warning(f"Stamping file name failed for {file} to {outf}")
+
+
+def safe_name(file_name: str):
+    return re.sub(pattern="[-!@#$%^&*()\s]", repl="_", string=file_name)

telewater/watermark/__init__.py

@@ -0,0 +1,4 @@
+"""Watermarking"""
+
+from .image import watermark_image
+from .video import watermark_video

telewater/watermark/image.py

@@ -0,0 +1,44 @@
+import logging
+
+from PIL import Image
+
+from telewater import conf
+from telewater.utils import cleanup
+
+
+def to_png(file: str) -> str:
+    im = Image.open(file)
+    if not file.endswith(".png"):
+        png_file = f"{file}.png"
+        im.save(png_file)
+    else:
+        png_file = file
+
+    return png_file
+
+
+def _open_image(file: str):
+    im = Image.open(file)
+    logging.info(
+        f"""Opened Image {file}
+        Image format {im.format}
+        Mode {im.mode}
+        Size {im.size}
+            Height {im.height}
+            Width {im.width}
+    """
+    )
+    im = im.convert("RGBA")
+
+    return im
+
+
+def watermark_image(image_file: str, watermark_file: str = "image.png"):
+    png = to_png(image_file)
+    im = _open_image(png)
+    wt = _open_image(watermark_file)
+    im.alpha_composite(im=wt, dest=(conf.config.x_off, conf.config.y_off))
+    output_file = f"watered_{png}"
+    im.save(output_file)
+    cleanup(png)
+    return output_file

telewater/watermark/video.py

@@ -0,0 +1,22 @@
+""" Watermark to videos is applied here."""
+
+import logging
+import os
+
+from telewater import conf
+
+
+def watermark_video(video_file: str, watermark_file: str = "image.png"):
+    """Apply watermark to video or gifs"""
+    output_file = f"watered_{video_file}"
+    command = f'ffmpeg -i {video_file} \
+        -i {watermark_file} \
+        -an -dn -sn -r {conf.config.frame_rate} \
+        -preset {conf.config.preset} \
+        -tune zerolatency  -tune fastdecode \
+        -filter_complex "overlay={conf.config.x_off}:{conf.config.y_off}" \
+        {output_file}'
+
+    logging.info(f"Running command {command}")
+    os.system(command)
+    return output_file

test.py

@@ -17,6 +17,7 @@
 
 class TestData(BaseModel):
     """This class defines the schema of the test.yml file"""
+
     bots: List[str]
     video_file: str
 
@@ -31,7 +32,7 @@ class TestData(BaseModel):
 
 
 async def general_test():
-    async with TelegramClient('telwater_user', API_ID, API_HASH) as client:
+    async with TelegramClient("telwater_user", API_ID, API_HASH) as client:
         me = await client.get_me()
         print(f"Logged in as {me.first_name}")
 
@@ -48,5 +49,6 @@ async def general_test():
 
         print("Forwarded the message to saved messages")
 
-if __name__ == '__main__':
+
+if __name__ == "__main__":
     asyncio.run(general_test())