Autodownload files

tg/config.py

@@ -1,10 +1,11 @@
-import os
 import configparser
 import mailcap
 import mimetypes
+import os
 
 DEFAULT_CONFIG = os.path.expanduser("~/.config/tg/tg.conf")
 DEFAULT_FILES = os.path.expanduser("~/.cache/tg/")
+max_download_size = "10MB"
 
 
 def get_cfg(config=DEFAULT_CONFIG):

tg/controllers/__init__.py

@@ -1,14 +1,14 @@
 import logging
-import os
 import threading
 from tempfile import NamedTemporaryFile
 
 from telegram.client import Telegram
 
-from tg.utils import notify, handle_exception, suspend
+from tg import config
 from tg.models import Model
-from tg.views import View
 from tg.msg import MsgProxy
+from tg.utils import handle_exception, notify, suspend
+from tg.views import View
 
 log = logging.getLogger(__name__)
 
@@ -44,10 +44,13 @@ class Controller:
         log.debug("Downloading msg: %s", msg.msg)
         file_id = msg.file_id
         if file_id:
-            log.info("Downloading file: file_id=%s", file_id)
-            self.model.downloads[file_id] = (msg["chat_id"], msg["id"])
-            self.tg.download_file(file_id=file_id)
-            log.info("Downloaded: file_id=%s", file_id)
+            self.download(file_id, msg["chat_id"], msg["id"])
+
+    def download(self, file_id: int, chat_id, msg_id):
+        log.info("Downloading file: file_id=%s", file_id)
+        self.model.downloads[file_id] = (chat_id, msg_id)
+        self.tg.download_file(file_id=file_id)
+        log.info("Downloaded: file_id=%s", file_id)
 
     def open_current_msg(self):
         msg = MsgProxy(self.model.current_msg())
@@ -68,10 +71,8 @@ class Controller:
                 s.run(path)
 
     def handle_msgs(self) -> str:
-        # set width to 0.25, move window to left
-        # refresh everything
         self.view.chats.resize(0.2)
-        self.view.msgs.resize(0.2)
+        self.view.msgs.resize(0.8)
         self.refresh_chats()
 
         while True:
@@ -139,8 +140,6 @@ class Controller:
                 return "BACK"
 
     def handle_chats(self) -> None:
-        # set width to 0.5, move window to center?
-        # refresh everything
         self.view.chats.resize(0.5)
         self.view.msgs.resize(0.5)
         self.refresh_chats()
@@ -198,10 +197,12 @@ class Controller:
 
     @handle_exception
     def update_new_msg(self, update):
-        msg = update["message"]
+        msg = MsgProxy(update["message"])
         chat_id = msg["chat_id"]
         self.model.msgs.add_message(chat_id, msg)
         self.refresh_msgs()
+        if msg.file_id and msg.size <= config.max_download_size:
+            self.download(msg.file_id, chat_id, msg['id'])
 
         # notify
         user_id = msg["sender_user_id"]

tg/main.py

@@ -1,16 +1,15 @@
 import logging
 import logging.handlers
 import threading
-from curses import wrapper, window
+from curses import window, wrapper
 from functools import partial
 
 from telegram.client import Telegram
 
+from tg import config, utils
 from tg.controllers import Controller
 from tg.models import Model
 from tg.views import View
-from tg import config, utils
-
 
 log = logging.getLogger(__name__)
 
@@ -59,6 +58,9 @@ def main():
         tdlib_verbosity=cfg.get("tdlib_verbosity", 0),
         library_path=cfg.get("library_path"),
     )
+    config.max_download_size = utils.parse_size(
+        cfg.get("max_download_size", config.max_download_size)
+    )
     tg.login()
     wrapper(partial(run, tg))
 

tg/msg.py

@@ -1,4 +1,5 @@
 import logging
+
 from tg import utils
 
 log = logging.getLogger(__name__)
@@ -60,6 +61,11 @@ class MsgProxy:
 
     @property
     def size(self):
+        doc = self.get_doc(self.msg)
+        return doc["size"]
+
+    @property
+    def human_size(self):
         doc = self.get_doc(self.msg)
         return utils.humanize_size(doc["size"])
 

tg/utils.py

@@ -11,6 +11,16 @@ from tg import config
 
 log = logging.getLogger(__name__)
 
+units = {"B": 1, "KB": 10**3, "MB": 10**6, "GB": 10**9, "TB": 10**12}
+
+
+def parse_size(size):
+    if size[-2].isalpha():
+        number, unit = size[:-2], size[-2:]
+    else:
+        number, unit = size[:-1], size[-1:]
+    return int(float(number)*units[unit])
+
 
 def humanize_size(
     num, suffix="B", suffixes=("", "K", "M", "G", "T", "P", "E", "Z")

tg/views/__init__.py

@@ -304,7 +304,7 @@ def parse_content(content: Dict[str, Any]) -> str:
     fields = dict(
         name=msg.file_name,
         duration=msg.duration,
-        size=msg.size,
+        size=msg.human_size,
         download=get_download(msg.local, msg.size),
     )
     info = ", ".join(f"{k}={v}" for k, v in fields.items() if v)