vid_utils.py 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103
  1. import re
  2. import os
  3. from glob import glob, escape
  4. from subprocess import Popen, PIPE
  5. from time import strftime, strptime, sleep
  6. from contextlib import contextmanager
  7. from telegram import InlineKeyboardButton
  8. class BadLink(Exception):
  9. pass
  10. class Video:
  11. def __init__(self, link, init_keyboard=False):
  12. self.link = link
  13. self.file_name = None
  14. if init_keyboard:
  15. self.formats = self.get_formats()
  16. self.keyboard = self.generate_keyboard()
  17. def get_formats(self):
  18. formats = []
  19. cmd = "youtube-dl -F {}".format(self.link)
  20. p = Popen(cmd, shell=True, stdout=PIPE, stderr=PIPE).communicate()
  21. it = iter(str(p[0], 'utf-8').split('\n')) # iterator of output lines
  22. try:
  23. while "code extension" not in next(it): pass # Remove garbage lines
  24. except StopIteration:
  25. raise BadLink # Isn't a valid youtube link
  26. while True:
  27. try:
  28. line = next(it)
  29. if not line:
  30. raise StopIteration # Usually the last line is empty
  31. if "video only" in line:
  32. continue # I don't need video without audio
  33. except StopIteration:
  34. break
  35. else:
  36. format_code, extension, resolution, *_ = line.strip().split()
  37. formats.append([format_code, extension, resolution])
  38. return formats
  39. def generate_keyboard(self):
  40. """ Generate a list of InlineKeyboardButton of resolutions """
  41. kb = []
  42. for code, extension, resolution in self.formats:
  43. kb.append([InlineKeyboardButton("{0}, {1}".format(extension, resolution),
  44. callback_data="{} {}".format(code, self.link))]) # maybe callback_data can support a list or tuple?
  45. return kb
  46. def download(self, resolution_code):
  47. cmd = "youtube-dl -f {0} {1}".format(resolution_code, self.link)
  48. p = Popen(cmd, shell=True, stdout=PIPE, stderr=PIPE).communicate()
  49. for line in str(p[0], 'utf-8').split('\n'):
  50. if "[download] Destination:" in line:
  51. self.file_name = line[24:] # name of the file
  52. def check_dimension(self):
  53. if os.path.getsize(self.file_name) > 50 * 1024 * 1023:
  54. os.system('split -b 49M "{0}" "{1}"'.format(self.file_name, self.file_name))
  55. os.remove(self.file_name)
  56. return glob(escape(self.file_name) + '*')
  57. @contextmanager
  58. def send(self):
  59. files = self.check_dimension() # split if size >= 50MB
  60. yield files
  61. for f in files: #removing old files
  62. os.remove(f)
  63. #__________________________OLD STUFFS, TOUCH CAREFULLY__________________________
  64. # this is the soft-split version, require avconv, but the audio isn't synchronized, avconv's problems :(
  65. '''
  66. def get_duration(filepath): # get duration in seconds
  67. cmd = "avconv -i %s" % filepath
  68. p = Popen(cmd, shell=True, stdout=PIPE, stderr=PIPE)
  69. di = p.communicate()
  70. for line in di:
  71. if line.rfind(b"Duration") > 0:
  72. duration = str(re.findall(b"Duration: (\d+:\d+:[\d.]+)", line)[0])
  73. return 3600 * int(duration[2: 4]) + 60 * int(duration[5: 7]) + int(duration[8: 10])
  74. def check_dimension(f): # if f is bigger than 50MB split it in subvideos
  75. if os.path.getsize(f) > 50 * 1024 * 1023:
  76. duration = get_duration(f)
  77. for i in range(0, duration, 180):
  78. start = strftime("%H:%M:%S", strptime('{0} {1} {2}'.format(i // 3600, (i // 60) % 60, i % 60), "%H %M %S")) # TODO this is not pythonic code!
  79. os.system("""avconv -i '{0}' -vcodec copy -acodec copy -ss {1} -t {2} 'part_{3}.mp4'""".format(f, start, 180, (i // 180) % 180))
  80. os.remove(f) # delete original file
  81. '''