refactor output to its own function

Almost the same code for all the fetchers.

lib/svtplay_dl/fetcher/hds.py

@@ -1,16 +1,13 @@
 # ex:ts=4:sw=4:sts=4:et
 # -*- tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*-
 from __future__ import absolute_import
-import sys
 import base64
-import re
 import struct
 import logging
 import binascii
-import os
 import xml.etree.ElementTree as ET
 
-from svtplay_dl.output import progressbar, progress_stream, ETA
+from svtplay_dl.output import progressbar, progress_stream, ETA, output
 from svtplay_dl.utils import get_http_data, is_py2_old, is_py2, is_py3
 from svtplay_dl.utils.urllib import urlparse
 from svtplay_dl.error import UIException
@@ -79,20 +76,9 @@ class HDS(VideoRetriever):
             antal = readbox(bootstrap, box[0])
         baseurl = self.kwargs["manifest"][0:self.kwargs["manifest"].rfind("/")]
 
-        if self.options.output != "-":
-            extension = re.search(r"(\.[a-z0-9]+)$", self.options.output)
-            if not extension:
-                self.options.output = "%s.flv" % self.options.output
-            log.info("Outfile: %s", self.options.output)
-            if os.path.isfile(self.options.output) and not self.options.force:
-                log.info("File already exists. use --force to overwrite")
-                return
-            file_d = open(self.options.output, "wb")
-        else:
-            if is_py3:
-                file_d = sys.stdout.buffer
-            else:
-                file_d = sys.stdout
+        file_d = output(self.options, self.options.output, "flv")
+        if file_d is None:
+            return
 
         metasize = struct.pack(">L", len(base64.b64decode(self.kwargs["metadata"])))[1:]
         file_d.write(binascii.a2b_hex(b"464c560105000000090000000012"))

lib/svtplay_dl/fetcher/hls.py

@@ -5,8 +5,8 @@ import sys
 import os
 import re
 
-from svtplay_dl.utils import get_http_data, is_py3
-from svtplay_dl.output import progressbar, progress_stream, ETA
+from svtplay_dl.utils import get_http_data
+from svtplay_dl.output import progressbar, progress_stream, ETA, output
 from svtplay_dl.log import log
 from svtplay_dl.utils.urllib import urlparse
 from svtplay_dl.error import UIException
@@ -80,20 +80,10 @@ class HLS(VideoRetriever):
             key = get_http_data(match.group(1))
             rand = os.urandom(16)
             decryptor = AES.new(key, AES.MODE_CBC, rand)
-        if self.options.output != "-":
-            extension = re.search(r"(\.[a-z0-9]+)$", self.options.output)
-            if not extension:
-                self.options.output = "%s.ts" % self.options.output
-            log.info("Outfile: %s", self.options.output)
-            if os.path.isfile(self.options.output) and not self.options.force:
-                log.info("File already exists. use --force to overwrite")
-                return
-            file_d = open(self.options.output, "wb")
-        else:
-            if is_py3:
-                file_d = sys.stdout.buffer
-            else:
-                file_d = sys.stdout
+
+        file_d = output(self.options, self.options.output, "ts")
+        if file_d is None:
+            return
 
         n = 0
         eta = ETA(len(files))

lib/svtplay_dl/fetcher/http.py

@@ -3,14 +3,11 @@
 from __future__ import absolute_import
 import sys
 import time
-import re
-import os
 
-from svtplay_dl.output import progress # FIXME use progressbar() instead
+from svtplay_dl.output import progress, output # FIXME use progressbar() instead
 from svtplay_dl.log import log
 from svtplay_dl.utils.urllib import urlopen, Request, HTTPError
 from svtplay_dl.fetcher import VideoRetriever
-from svtplay_dl.utils import is_py3
 
 class HTTP(VideoRetriever):
     def name(self):
@@ -32,22 +29,10 @@ class HTTP(VideoRetriever):
             total_size = 0
         total_size = int(total_size)
         bytes_so_far = 0
-        if self.options.output != "-":
-            extension = re.search(r"(\.[a-z0-9]+)$", self.url)
-            if extension:
-                self.options.output = self.options.output + extension.group(1)
-            else:
-                self.options.output = "%s.mp4" % self.options.output
-            log.info("Outfile: %s", self.options.output)
-            if os.path.isfile(self.options.output) and not self.options.force:
-                log.info("File already exists. use --force to overwrite")
-                return
-            file_d = open(self.options.output, "wb")
-        else:
-            if is_py3:
-                file_d = sys.stdout.buffer
-            else:
-                file_d = sys.stdout
+
+        file_d = output(self.options, self.url, "mp4")
+        if file_d is None:
+            return
 
         lastprogress = 0
         while 1:

lib/svtplay_dl/fetcher/rtmp.py

@@ -9,6 +9,7 @@ import os
 from svtplay_dl.log import log
 from svtplay_dl.utils import is_py2
 from svtplay_dl.fetcher import VideoRetriever
+from svtplay_dl.output import output
 
 class RTMP(VideoRetriever):
     def name(self):
@@ -23,17 +24,10 @@ class RTMP(VideoRetriever):
         if self.options.resume:
             args.append("-e")
 
-        extension = re.search(r"(\.[a-z0-9]+)$", self.url)
-        if self.options.output != "-":
-            if not extension:
-                self.options.output = "%s.flv" % self.options.output
-            else:
-                self.options.output = self.options.output + extension.group(1)
-            log.info("Outfile: %s", self.options.output)
-            if os.path.isfile(self.options.output) and not self.options.force:
-                log.info("File already exists. use --force to overwrite")
-                return
-            args += ["-o", self.options.output]
+        fild_d = output(self.options, self.options.output, "flv", False)
+        if fild_d is None:
+            return
+        args += ["-o", self.options.output]
         if self.options.silent or self.options.output == "-":
             args.append("-q")
         if self.options.other:

lib/svtplay_dl/output.py

@@ -3,9 +3,13 @@
 from __future__ import absolute_import
 import sys
 import time
+import re
+import os
 from datetime import timedelta
 
+from svtplay_dl.utils import is_py3
 from svtplay_dl.utils.terminal import get_terminal_size
+from svtplay_dl.log import log
 
 progress_stream = sys.stderr
 
@@ -108,3 +112,23 @@ def progressbar(total, pos, msg=""):
 
     progress_stream.write(fmt % (pos, total, bar, msg))
 
+def output(options, filename, extention="mp4", openfd=True):
+    file_d = -1
+    if options.output != "-":
+        ext = re.search(r"(\.[a-z0-9]+)$", filename)
+        if not ext:
+            options.output = "%s.%s" % (options.output, extention)
+        log.info("Outfile: %s", options.output)
+        if os.path.isfile(options.output) and not options.force:
+            log.info("File already exists. use --force to overwrite")
+            return None
+        if openfd:
+            file_d = open(options.output, "wb")
+    else:
+        if openfd:
+            if is_py3:
+                file_d = sys.stdout.buffer
+            else:
+                file_d = sys.stdout
+
+    return file_d