Pass user-agent to YoutubeDlCoprocessor

ArchiveTeam · Jan 29, 2015 · 6be81ab · 6be81ab
1 parent e8ffe6a
commit 6be81ab
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 4 deletions.
diff --git a/wpull/builder.py b/wpull/builder.py
@@ -1234,7 +1234,8 @@ def _build_youtube_dl_coprocessor(self, proxy_port):
             'YoutubeDlCoprocessor',
             self._args.youtube_dl_exe,
             (self._args.proxy_server_address, proxy_port),
-            self._args.directory_prefix,
+            root_path=self._args.directory_prefix,
+            user_agent = self._args.user_agent or self.default_user_agent
         )
 
         return coprocessor

diff --git a/wpull/coprocessor/youtubedl.py b/wpull/coprocessor/youtubedl.py
@@ -17,10 +17,12 @@
 
 
 class YoutubeDlCoprocessor(object):
-    def __init__(self, youtube_dl_path, proxy_address, root_path='.'):
+    def __init__(self, youtube_dl_path, proxy_address, root_path='.',
+                 user_agent=None):
         self._youtube_dl_path = youtube_dl_path
         self._proxy_address = proxy_address
         self._root_path = root_path
+        self._user_agent = user_agent
 
     @trollius.coroutine
     def process(self, url_item, request, response, file_writer_session):
@@ -32,7 +34,7 @@ def process(self, url_item, request, response, file_writer_session):
 
         session = Session(
             self._proxy_address, self._youtube_dl_path, self._root_path,
-            url_item, file_writer_session
+            url_item, file_writer_session, self._user_agent
         )
 
         url = url_item.url_info.url
@@ -46,12 +48,13 @@ def process(self, url_item, request, response, file_writer_session):
 
 class Session(object):
     def __init__(self, proxy_address, youtube_dl_path, root_path, url_item,
-                 file_writer_session):
+                 file_writer_session, user_agent):
         self._proxy_address = proxy_address
         self._youtube_dl_path = youtube_dl_path
         self._root_path = root_path
         self._url_item = url_item
         self._file_writer_session = file_writer_session
+        self._user_agent = user_agent
         self._temp_dir = None
 
     @trollius.coroutine
@@ -72,6 +75,9 @@ def run(self):
             url
         ]
 
+        if self._user_agent:
+            args.extend(['--user-agent', self._user_agent])
+
         youtube_dl_process = Process(
             args,
             stderr_callback=self._stderr_callback,