From bf1f60eca590e0348b2be22483a697596ce65ecf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1=C5=A1=20Ml=C3=A1dek?= <t@mldk.cz>
Date: Wed, 25 Apr 2018 14:21:48 +0200
Subject: [PATCH] grande refactor

---
 delojza.py | 102 +++++++++++++++++++++++++----------------------------
 1 file changed, 48 insertions(+), 54 deletions(-)

diff --git a/delojza.py b/delojza.py
index a9a9b55..e575dc9 100755
--- a/delojza.py
+++ b/delojza.py
@@ -46,7 +46,7 @@ def ytdl_has(url):
     return False
 
 
-def download_ydl(urls, subdir, date, extract=False):
+def download_ydl(urls, subdir, date, extract=False, filename=None):
     ydl_opts = {
         'noplaylist': True,
         'restrictfilenames': True,
@@ -69,21 +69,21 @@ def download_ydl(urls, subdir, date, extract=False):
                 shutil.move(globbed, out_dir)
 
 
-def download_raw(url, filename, subdir, date):
-    local_filename = f"{OUT_DIR}/{subdir}/" + "%s__%s" % (datestr(date), filename or url.split('/')[-1])
-    # local_filename = OUT_DIR + '/' + ("%s/" % subdir) if subdir else '' + datestr(date) + '__' + url.split('/')[-1]
-    r = requests.get(url, stream=True)
-    with open(local_filename, 'wb') as f:
-        for chunk in r.iter_content(chunk_size=1024):
-            if chunk:
-                f.write(chunk)
-    if not re.match(r'.*\..{3,5}$', os.path.split(local_filename)[-1]):
-        kind = filetype.guess(local_filename)
-        if kind is None:
-            logger.error("File has no extension and could not be determined!")
-        else:
-            logger.info('Moving file without extension... %s?' % kind.extension)
-            shutil.move(local_filename, local_filename + '.' + kind.extension)
+def download_raw(urls, subdir, date, extract=False, filename=None):
+    for url in urls:
+        local_filename = f"{OUT_DIR}/{subdir}/" + "%s__%s" % (datestr(date), filename or url.split('/')[-1])
+        r = requests.get(url, stream=True)
+        with open(local_filename, 'wb') as f:
+            for chunk in r.iter_content(chunk_size=1024):
+                if chunk:
+                    f.write(chunk)
+        if not re.match(r'.*\..{3,5}$', os.path.split(local_filename)[-1]):
+            kind = filetype.guess(local_filename)
+            if kind is None:
+                logger.error("File has no extension and could not be determined!")
+            else:
+                logger.info('Moving file without extension... %s?' % kind.extension)
+                shutil.move(local_filename, local_filename + '.' + kind.extension)
 
 
 def get_first_hashtag(message):
@@ -100,18 +100,15 @@ def get_first_hashtag(message):
 
 
 # noinspection PyBroadException
-def handle_url(bot, update):
-    ytdl_urls = list(filter(ytdl_has,
-                            map(lambda e: update.message.parse_entity(e),
-                                filter(lambda e: e.type == 'url',
-                                       update.message.entities))))
-    if len(ytdl_urls) > 0:
+def handle(urls, message, download, filename=None):
+    tries = 0
+    while tries < 3:
         try:
-            logger.info("Downloading %s" % ytdl_urls)
-            hashtag = get_first_hashtag(update.message)
+            logger.info("Downloading %s" % urls)
+            hashtag = get_first_hashtag(message)
 
             if hashtag == 'IGNORE':
-                update.message.reply_text('Ignoring...')
+                message.reply_text('Ignoring...')
                 return
 
             reply = 'Downloading'
@@ -120,13 +117,32 @@ def handle_url(bot, update):
                 reply += f' to "{hashtag}"'
             reply += '...'
 
-            if hashtag == 'AUDIO':
+            if hashtag == 'AUDIO' and download != download_raw:
                 reply += ' (And also guessing you want to extract the audio)'
-            update.message.reply_text(reply)
-            download_ydl(ytdl_urls, hashtag or '.', update.message.date, extract=hashtag == 'AUDIO')
+            message.reply_text(reply)
+            download(urls,
+                     hashtag or '.', message.date,
+                     extract=(hashtag == 'AUDIO'),
+                     filename=filename)
+            break
         except:
-            type, value, _ = sys.exc_info()
-            update.message.reply_text("Something is FUCKED: %s, %s" % (type, value))
+            _, exc_value, __ = sys.exc_info()
+            if "Timed out" not in str(exc_value):
+                message.reply_text("Something is FUCKED: %s" % exc_value)
+                break
+            else:
+                tries += 1
+    if tries == 3:
+        message.reply_text("Something is FUCKED, tried thrice and it's all wrong")
+
+
+def handle_url(bot, update):
+    ytdl_urls = list(filter(ytdl_has,
+                            map(lambda e: update.message.parse_entity(e),
+                                filter(lambda e: e.type == 'url',
+                                       update.message.entities))))
+    if len(ytdl_urls) > 0:
+        handle(ytdl_urls, update.message, download_ydl)
 
 
 # noinspection PyBroadException
@@ -146,30 +162,8 @@ def handle_rest(bot, update):
     elif update.message.voice is not None:
         file = update.message.voice.file_id
     if file is not None:
-        try:
-            url = bot.getFile(file).file_path
-            logger.info("Downloading %s" % url)
-            hashtag = get_first_hashtag(update.message)
-            if hashtag == 'IGNORE':
-                update.message.reply_text('Ignoring...')
-                return
-
-            if not hashtag and update.message.audio is not None:
-                hashtag = 'AUDIO'
-
-            reply = 'Downloading'
-            if hashtag:
-                mkdir_p(f'{OUT_DIR}/{hashtag}')
-                reply += f' to "{hashtag}"'
-            reply += '...'
-            update.message.reply_text(reply)
-
-            if hashtag:
-                mkdir_p(f'{OUT_DIR}/{hashtag}')
-            download_raw(url, filename, hashtag or '.', update.message.date)
-        except:
-            type, value, _ = sys.exc_info()
-            update.message.reply_text("Something is FUCKED: %s, %s" % (type, value))
+        url = bot.getFile(file).file_path
+        handle([url], update.message, download_raw, filename=filename)
 
 
 def start(bot, update):