You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

254 lines
7.0 KiB

  1. #! /usr/bin/env python
  2. import subprocess
  3. from .Utils import (
  4. get_encoding,
  5. get_filename,
  6. encode_list,
  7. os_type
  8. )
  9. class DownloadObject(object):
  10. '''
  11. Download videos using youtube-dl & subprocess.
  12. Params
  13. youtubedl_path: Absolute path of youtube-dl.
  14. data_hook: Can be any function with one parameter, the data.
  15. logger: Can be any logger which implements log().
  16. Accessible Methods
  17. download()
  18. Params: URL to download
  19. Options list e.g. ['--help']
  20. Return: DownlaodObject.OK
  21. DownloadObject.ERROR
  22. DownloadObject.STOPPED
  23. DownloadObject.ALREADY
  24. stop()
  25. Params: None
  26. Acessible Variables
  27. files_list: Python list that contains all the files DownloadObject
  28. instance has downloaded.
  29. Data_hook Keys
  30. 'playlist_index',
  31. 'playlist_size',
  32. 'filesize',
  33. 'filename',
  34. 'percent',
  35. 'status',
  36. 'speed',
  37. 'eta'
  38. '''
  39. # download() return codes
  40. OK = 0
  41. ERROR = 1
  42. STOPPED = 2
  43. ALREADY = 3
  44. def __init__(self, youtubedl_path, data_hook=None, logger=None):
  45. self.youtubedl_path = youtubedl_path
  46. self.data_hook = data_hook
  47. self.logger = logger
  48. self.files_list = []
  49. self._return_code = 0
  50. self._proc = None
  51. self._init_data()
  52. def _init_data(self):
  53. ''' Keep the __init__() clean. '''
  54. self._data = {
  55. 'playlist_index': None,
  56. 'playlist_size': None,
  57. 'filesize': None,
  58. 'filename': None,
  59. 'percent': None,
  60. 'status': None,
  61. 'speed': None,
  62. 'eta': None
  63. }
  64. def download(self, url, options):
  65. self._return_code = self.OK
  66. cmd = self._get_cmd(url, options)
  67. cmd = self._encode_cmd(cmd)
  68. info = self._get_process_info()
  69. self._proc = self._create_process(cmd, info)
  70. while self._proc_is_alive():
  71. stdout, stderr = self._read()
  72. data = extract_data(stdout)
  73. updated = self._update_data(data)
  74. if stderr != '':
  75. self._return_code = self.ERROR
  76. self._log(stderr)
  77. if updated:
  78. self._hook_data()
  79. return self._return_code
  80. def stop(self):
  81. if self._proc is not None:
  82. self._proc.kill()
  83. self._return_code = self.STOPPED
  84. def _update_data(self, data):
  85. ''' Update self._data from data.
  86. Return True if updated else return False.
  87. '''
  88. updated = False
  89. for key in data:
  90. if key == 'filename':
  91. # Save full file path on files_list
  92. self._add_on_files_list(data['filename'])
  93. # Keep only the filename not the path on data['filename']
  94. data['filename'] = get_filename(data['filename'])
  95. if key == 'status':
  96. # Set self._return_code to already downloaded
  97. if data[key] == 'already_downloaded':
  98. self._return_code = self.ALREADY
  99. # Trash that key
  100. data[key] = None
  101. self._data[key] = data[key]
  102. updated = True
  103. return updated
  104. def _add_on_files_list(self, filename):
  105. self.files_list.append(filename)
  106. def _log(self, data):
  107. if self.logger is not None:
  108. self.logger.log(data)
  109. def _hook_data(self):
  110. ''' Pass self._data back to data_hook. '''
  111. if self.data_hook is not None:
  112. self.data_hook(self._data)
  113. def _proc_is_alive(self):
  114. ''' Return True if self._proc is alive. '''
  115. if self._proc is None:
  116. return False
  117. return self._proc.poll() is None
  118. def _read(self):
  119. ''' Read subprocess stdout, stderr. '''
  120. stdout = self._read_stdout()
  121. if stdout == '':
  122. stderr = self._read_stderr()
  123. else:
  124. stderr = ''
  125. return stdout, stderr
  126. def _read_stdout(self):
  127. if self._proc is None:
  128. return ''
  129. stdout = self._proc.stdout.readline()
  130. return stdout.rstrip()
  131. def _read_stderr(self):
  132. if self._proc is None:
  133. return ''
  134. stderr = self._proc.stderr.readline()
  135. return stderr.rstrip()
  136. def _create_process(self, cmd, info):
  137. return subprocess.Popen(cmd,
  138. stdout=subprocess.PIPE,
  139. stderr=subprocess.PIPE,
  140. startupinfo=info)
  141. def _get_cmd(self, url, options):
  142. ''' Return command for subprocess. '''
  143. if os_type == 'nt':
  144. cmd = [self.youtubedl_path] + options + [url]
  145. else:
  146. cmd = ['python', self.youtubedl_path] + options + [url]
  147. return cmd
  148. def _encode_cmd(self, cmd):
  149. ''' Encode command for subprocess.
  150. Refer to http://stackoverflow.com/a/9951851/35070
  151. '''
  152. encoding = get_encoding()
  153. if encoding is not None:
  154. cmd = encode_list(cmd, encoding)
  155. return cmd
  156. def _get_process_info(self):
  157. ''' Hide subprocess window on Windows. '''
  158. if os_type == 'nt':
  159. info = subprocess.STARTUPINFO()
  160. info.dwFlags |= subprocess.STARTF_USESHOWWINDOW
  161. return info
  162. else:
  163. return None
  164. def extract_data(stdout):
  165. ''' Extract data from youtube-dl stdout. '''
  166. data_dictionary = {}
  167. stdout = [s for s in stdout.split(' ') if s != '']
  168. if len(stdout) == 0:
  169. return data_dictionary
  170. header = stdout.pop(0)
  171. if header[0] == '[' and header[-1] == ']':
  172. header = header.replace('[', '').replace(']', '')
  173. if header == 'download':
  174. data_dictionary['status'] = 'download'
  175. # Get filename
  176. if stdout[0] == 'Destination:':
  177. data_dictionary['filename'] = ' '.join(stdout[1:])
  178. # Get progress info
  179. if '%' in stdout[0]:
  180. if stdout[0] == '100%':
  181. data_dictionary['speed'] = ''
  182. data_dictionary['eta'] = ''
  183. else:
  184. data_dictionary['percent'] = stdout[0]
  185. data_dictionary['filesize'] = stdout[2]
  186. data_dictionary['speed'] = stdout[4]
  187. data_dictionary['eta'] = stdout[6]
  188. # Get playlist info
  189. if stdout[0] == 'Downloading' and stdout[1] == 'video':
  190. data_dictionary['playlist_index'] = stdout[2]
  191. data_dictionary['playlist_size'] = stdout[4]
  192. # Get file already downloaded status
  193. if stdout[-1] == 'downloaded':
  194. data_dictionary['status'] = 'already_downloaded'
  195. elif header == 'ffmpeg':
  196. data_dictionary['status'] = 'post_process'
  197. else:
  198. data_dictionary['status'] = 'pre_process'
  199. return data_dictionary