diff --git a/robot.py b/robot.py index 95971cf..966ba39 100755 --- a/robot.py +++ b/robot.py @@ -1,4 +1,5 @@ #!/usr/bin/env python3 +import errno import logging import os import re @@ -21,6 +22,20 @@ logging.basicConfig(level=logging.INFO, logger = logging.getLogger("kunsax") +def datestr(date): + return date.strftime("%Y-%m-%d@%H%M") + + +def mkdir_p(path): + try: + os.makedirs(path) + except OSError as exc: # Python >2.5 + if exc.errno == errno.EEXIST and os.path.isdir(path): + pass + else: + raise + + def ytdl_has(url): ies = youtube_dl.extractor.gen_extractors() for ie in ies: @@ -31,28 +46,26 @@ def ytdl_has(url): return False -def datestr(date): - return date.strftime("%Y-%m-%d@%H%M") - - -def download_ydl(urls, date): +def download_ydl(urls, subdir, date): ydl_opts = { 'noplaylist': True, 'restrictfilenames': True, 'download_archive': DIR + '/downloaded.lst', - 'outtmpl': TMP_DIR + '/' + datestr(date) + '__%(title)s__%(id)s.%(ext)s' + 'outtmpl': f'{TMP_DIR}/' + datestr(date) + '__%(title)s__%(id)s.%(ext)s' } with youtube_dl.YoutubeDL(ydl_opts) as ydl: ydl.download(urls) + out_dir = f'{OUT_DIR}/{subdir}/' for filename in map(ydl.prepare_filename, map(ydl.extract_info, urls)): globbeds = glob(os.path.splitext(filename)[0] + '.*') for globbed in globbeds: - logger.info("Moving %s to %s..." % (globbed, OUT_DIR)) - shutil.move(globbed, OUT_DIR) + logger.info("Moving %s to %s..." % (globbed, out_dir)) + shutil.move(globbed, out_dir) -def download_raw(url, date): - local_filename = OUT_DIR + '/' + datestr(date) + '__' + url.split('/')[-1] +def download_raw(url, subdir, date): + local_filename = f"{OUT_DIR}/{subdir}/" + "%s__%s" % (datestr(date), 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): @@ -67,17 +80,28 @@ def download_raw(url, date): shutil.move(local_filename, local_filename + '.' + kind.extension) +def get_first_hashtag(message): + hashtags = list(filter(lambda e: e.type == 'hashtag', message.entities)) + if len(hashtags) == 0: + return None + hashtag = message.parse_entity(hashtags[0]) + return hashtag[1:].upper() + + # noinspection PyBroadException def handle_url(bot, update): ytdl_urls = list(filter(ytdl_has, - map(lambda e: update.message.text[e.offset:e.offset + e.length], + map(lambda e: update.message.parse_entity(e), filter(lambda e: e.type == 'url', update.message.entities)))) if len(ytdl_urls) > 0: try: logger.info("Downloading %s" % ytdl_urls) - update.message.reply_text('Downloading now...') - download_ydl(ytdl_urls, update.message.date) + hashtag = get_first_hashtag(update.message) + if hashtag: + mkdir_p(f'{OUT_DIR}/{hashtag}') + update.message.reply_text('Downloading%s...' % f' to "{hashtag}"' if hashtag else '') + download_ydl(ytdl_urls, hashtag or '.', update.message.date) except: type, value, _ = sys.exc_info() update.message.reply_text("Something is FUCKED: %s, %s" % (type, value)) @@ -98,9 +122,12 @@ def handle_rest(bot, update): if file is not None: try: url = bot.getFile(file).file_path - update.message.reply_text('Downloading now...') - logger.info("Downloading '%s'" % url) - download_raw(url, update.message.date) + logger.info("Downloading %s" % url) + hashtag = get_first_hashtag(update.message) + if hashtag: + mkdir_p(f'{OUT_DIR}/{hashtag}') + update.message.reply_text('Downloading%s...' % f' to "{hashtag}"' if hashtag else '') + download_raw(url, hashtag or '.', update.message.date) except: type, value, _ = sys.exc_info() update.message.reply_text("Something is FUCKED: %s, %s" % (type, value))