Browse Source

first version

thezero 2 years ago
parent
commit
586668c7e0
4 changed files with 227 additions and 42 deletions
  1. 133 0
      .gitignore
  2. 36 0
      db.py
  3. 21 17
      main.py
  4. 37 25
      vid_utils.py

+ 133 - 0
.gitignore

@@ -0,0 +1,133 @@
+viddb.sqlite3
+out/
+conf/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/

+ 36 - 0
db.py

@@ -0,0 +1,36 @@
+import os
+import sqlite3
+
+class VidDatabase():
+
+	def __init__(self, path):
+		self.path = path
+		self.is_valid = os.path.isfile(self.path)
+		self.con = sqlite3.connect(self.path)
+
+	def __del__(self):
+		self.con.close()
+
+	def create(self):
+		cur = self.con.cursor()
+		cur.execute("CREATE TABLE vid (link TEXT NOT NULL, code TEXT NOT NULL);")
+		self.con.commit()
+
+	def select_vid(self, vid):
+		cur = self.con.cursor()
+		cur.execute("SELECT link, code FROM vid WHERE rowid = ?;", (vid,))
+		rows = cur.fetchall()
+		video = rows[0] if len(rows) > 0 else None
+		self.con.commit()
+		return video
+
+	def insert_vid(self, link, code):
+		cur = self.con.cursor()
+		cur.execute("INSERT INTO vid (link, code) VALUES (?, ?);", (link, code))
+		self.con.commit()
+		return cur.lastrowid
+
+	def delete_vid(self, link):
+		cur = self.con.cursor()
+		cur.execute("DELETE FROM vid WHERE link = ?;", link)
+		self.con.commit()

+ 21 - 17
main.py

@@ -9,35 +9,39 @@ logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s
 logger = logging.getLogger(__name__)
 
 
-def get_format(bot, update):
-    logger.info("from {}: {}".format(update.message.chat_id, update.message.text)) # "history"
+def get_format(update, context):
+    logger.info("from {}: {}".format(update.message.chat_id, update.message.text))  # "history"
 
     try:
-        video = Video(update.message.text, init_keyboard=True)
-    except BadLink:
-        update.message.reply_text("Bad link")
+        video = Video(link=update.message.text, init_keyboard=True)
+    except BadLink as e:
+        update.message.reply_text("Bad link: {}".format(e))
     else:
         reply_markup = InlineKeyboardMarkup(video.keyboard)
         update.message.reply_text('Choose format:', reply_markup=reply_markup)
 
 
-def download_choosen_format(bot, update):
+def download_choosen_format(update, context):
     query = update.callback_query
-    resolution_code, link = query.data.split(' ', 1)
-    
-    bot.edit_message_text(text="Downloading...",
-                          chat_id=query.message.chat_id,
-                          message_id=query.message.message_id)
-    
-    video = Video(link)
-    video.download(resolution_code)
-    
+    video_index = query.data
+
+    context.bot.edit_message_text(text="Downloading...",
+                                  chat_id=query.message.chat_id,
+                                  message_id=query.message.message_id)
+
+    video = Video(vid=video_index)
+    video.download()
+
     with video.send() as files:
         for f in files:
-            bot.send_document(chat_id=query.message.chat_id, document=open(f, 'rb'))
+            logger.log("Sending... {} ".format(f))
+            context.bot.send_document(chat_id=query.message.chat_id, document=open(f, 'rb'))
+
+    context.bot.delete_message(chat_id=query.message.chat_id,
+                               message_id=query.message.message_id)
 
 
-updater = Updater(token=YOUR_TOKEN)
+updater = Updater(token="")
 
 updater.dispatcher.add_handler(MessageHandler(Filters.text, get_format))
 updater.dispatcher.add_handler(CallbackQueryHandler(download_choosen_format))

+ 37 - 25
vid_utils.py

@@ -1,79 +1,91 @@
-import re
 import os
 from glob import glob, escape
 from subprocess import Popen, PIPE
-from time import strftime, strptime, sleep
 from contextlib import contextmanager
 
 from telegram import InlineKeyboardButton
 
+from db import VidDatabase
+
 class BadLink(Exception):
     pass
 
 
 class Video:
-    def __init__(self, link, init_keyboard=False):
-        self.link = link
-        self.file_name = None
-        
+    def __init__(self, link=None, vid=None, init_keyboard=False):
+        self.db = VidDatabase("viddb.sqlite3")
+        if not self.db.is_valid:
+            # Database file not present
+            # Create a new database
+            self.db.create()
+
+        if vid is None and link is not None:
+            self.link = link
+            self.file_name = None
+        elif vid is not None and link is None:
+            self.link, self.code = self.db.select_vid(vid)
+        else:
+            raise Exception('what is going on?')
+
         if init_keyboard:
             self.formats = self.get_formats()
             self.keyboard = self.generate_keyboard()
 
     def get_formats(self):
-        formats = []
+        formats = {}
 
-        cmd = "youtube-dl -F {}".format(self.link)
-        p = Popen(cmd, shell=True, stdout=PIPE, stderr=PIPE).communicate()
-        it = iter(str(p[0], 'utf-8').split('\n')) # iterator of output lines
+        p = Popen(["youtube-dl", "-F", self.link], stdout=PIPE, stderr=PIPE).communicate()
+        it = iter(str(p[0], 'utf-8').split('\n'))  # iterator of output lines
 
         try:
-            while "code  extension" not in next(it): pass # Remove garbage lines
+            while "code  extension" not in next(it):
+                pass  # Remove garbage lines
         except StopIteration:
-            raise BadLink # Isn't a valid youtube link
+            raise BadLink("youtube-dl couldn't download the link you provided")  # Isn't a valid youtube link
 
         while True:
             try:
                 line = next(it)
                 if not line:
-                    raise StopIteration # Usually the last line is empty
+                    raise StopIteration  # Usually the last line is empty
                 if "video only" in line:
-                    continue # I don't need video without audio
+                    continue  # I don't need video without audio
             except StopIteration:
                 break
             else:
                 format_code, extension, resolution, *_ = line.strip().split()
-                formats.append([format_code, extension, resolution])
+                key = '{},{}'.format(extension, resolution)
+                index = self.db.insert_vid(self.link, format_code)
+                formats[key] = index
         return formats
 
     def generate_keyboard(self):
         """ Generate a list of InlineKeyboardButton of resolutions """
         kb = []
 
-        for code, extension, resolution in self.formats:
-            kb.append([InlineKeyboardButton("{0}, {1}".format(extension, resolution),
-                                     callback_data="{} {}".format(code, self.link))]) # maybe callback_data can support a list or tuple?
+        for key in self.formats.keys():
+            cb = "{}".format(self.formats[key])
+            kb.append([InlineKeyboardButton(key, callback_data=cb)])
         return kb
 
-    def download(self, resolution_code):
-        cmd = "youtube-dl -f {0} {1}".format(resolution_code, self.link)
-        p = Popen(cmd, shell=True, stdout=PIPE, stderr=PIPE).communicate()
+    def download(self):
+        p = Popen(["youtube-dl", "-f", self.code, self.link], stdout=PIPE, stderr=PIPE).communicate()
 
         for line in str(p[0], 'utf-8').split('\n'):
             if "[download] Destination:" in line:
-                self.file_name = line[24:] # name of the file
+                self.file_name = line[24:]  # name of the file
 
     def check_dimension(self):
         if os.path.getsize(self.file_name) > 50 * 1024 * 1023:
-            os.system('split -b 49M "{0}" "{1}"'.format(self.file_name, self.file_name))
+            Popen(["split", "-b", "49M", self.file_name, self.file_name])
             os.remove(self.file_name)
         return glob(escape(self.file_name) + '*')
 
     @contextmanager
     def send(self):
-        files = self.check_dimension() # split if size >= 50MB
+        files = self.check_dimension()  # split if size >= 50MB
         yield files
-        for f in files: #removing old files
+        for f in files:  # removing old files
             os.remove(f)