Refactor stream input/output.

- Now uses pbs library to handle subprocesses.
- Removed -c option as not all plugins may be based on subprocesses.

10 files changed, 109 insertions, 80 deletions

diff --git a/README.md b/README.md
index a04b0aa..06fa8e2 100644
--- a/README.md
+++ b/README.md
@@ -59,14 +59,14 @@ Livestreamer is also a library. Short example:
     streams = channel.get_streams()
 
     stream = streams["720p"]
-    stream.open()
+    fd = stream.open()
 
     while True:
-        data = stream.read(1024)
+        data = fd.read(1024)
         if len(data) == 0:
             break
 
         # do something with data
 
-    stream.close()
+    fd.close()
 
diff --git a/setup.py b/setup.py
index 78552f8..44e7f72 100644
--- a/setup.py
+++ b/setup.py
@@ -4,7 +4,7 @@ from setuptools import setup, find_packages
 from sys import version_info
 
 version = "1.0.0"
-deps = []
+deps = ["pbs"]
 
 # require argparse on Python <2.7 and <3.2
 if (version_info[0] == 2 and version_info[1] < 7) or \
diff --git a/src/livestreamer/__init__.py b/src/livestreamer/__init__.py
index 64886a8..1934e2d 100644
--- a/src/livestreamer/__init__.py
+++ b/src/livestreamer/__init__.py
@@ -1,4 +1,5 @@
 from livestreamer import plugins
+from livestreamer import stream
 from livestreamer.compat import urlparse
 
 def resolve_url(url):
@@ -21,5 +22,6 @@ def get_plugins():
 PluginError = plugins.PluginError
 NoStreamsError = plugins.NoStreamsError
 NoPluginError = plugins.NoPluginError
+StreamError = stream.StreamError
 
 plugins.load_plugins(plugins)
diff --git a/src/livestreamer/cli.py b/src/livestreamer/cli.py
index d8d1bc0..75f645e 100644
--- a/src/livestreamer/cli.py
+++ b/src/livestreamer/cli.py
@@ -1,7 +1,8 @@
 #!/usr/bin/env python3
 
-import sys, os
+import sys, os, pbs
 import livestreamer
+from livestreamer.compat import input
 
 parser = livestreamer.utils.ArgumentParser(description="Util to play various livestreaming services in a custom player",
                                            fromfile_prefix_chars="@")
@@ -9,7 +10,6 @@ parser.add_argument("url", help="URL to stream", nargs="?")
 parser.add_argument("stream", help="stream to play", nargs="?")
 parser.add_argument("-p", "--player", metavar="player", help="commandline for player", default="vlc")
 parser.add_argument("-o", "--output", metavar="filename", help="write stream to file instead of playing it")
-parser.add_argument("-c", "--cmdline", action="store_true", help="print commandline used internally to play stream")
 parser.add_argument("-l", "--plugins", action="store_true", help="print installed plugins")
 
 RCFILE = os.path.expanduser("~/.livestreamerrc")
@@ -18,6 +18,33 @@ def exit(msg):
     sys.stderr.write("error: " + msg + "\n")
     sys.exit()
 
+def msg(msg):
+    sys.stderr.write(msg + "\n")
+
+def write_stream(fd, out, progress):
+    written = 0
+
+    while True:
+        data = fd.read(8192)
+        if len(data) == 0:
+            break
+
+        try:
+            out.write(data)
+        except IOError:
+            break
+
+        written += len(data)
+
+        if progress:
+            sys.stdout.write(("\rWritten {0} bytes").format(written))
+
+    if progress and written > 0:
+        sys.stdout.write("\n")
+
+    fd.close()
+    out.close()
+
 def handle_url(args):
     try:
         channel = livestreamer.resolve_url(args.url)
@@ -39,26 +66,47 @@ def handle_url(args):
     if args.stream:
         if args.stream in streams:
             stream = streams[args.stream]
-            cmdline = stream.cmdline(args.output or "-")
 
-            if args.cmdline:
-                print(cmdline.format())
-                sys.exit()
+            try:
+                fd = stream.open()
+            except livestreamer.StreamError as err:
+                exit(("Could not open stream - {0}").format(err))
+
+            progress = False
+
+            if args.output:
+                progress = True
+
+                if os.path.exists(args.output):
+                    answer = input(("File output {0} already exists! Overwrite it? [y/N] ").format(args.output))
+                    answer = answer.strip().lower()
+
+                    if answer != "y":
+                        sys.exit()
+
+                try:
+                    out = open(args.output, "wb")
+                except IOError as err:
+                    exit(("Failed to open file {0} - ").format(args.output, err))
             else:
-                if not args.output:
-                    cmdline.pipe = ("{0} -").format(args.player)
+                cmd = args.player + " -"
+                player = pbs.sh("-c", cmd, _bg=True, _out=sys.stdout, _err=sys.stderr)
+                out = player.process.stdin
 
-                os.system(cmdline.format())
+            try:
+                write_stream(fd, out, progress)
+            except KeyboardInterrupt:
+                sys.exit()
         else:
-            print(("This channel does not have stream: {0}").format(args.stream))
-            print(("Valid streams: {0}").format(validstreams))
+            msg(("This channel does not have stream: {0}").format(args.stream))
+            msg(("Valid streams: {0}").format(validstreams))
     else:
-        print(("Found streams: {0}").format(validstreams))
+        msg(("Found streams: {0}").format(validstreams))
 
 
 def print_plugins():
     pluginlist = list(livestreamer.get_plugins().keys())
-    print(("Installed plugins: {0}").format(", ".join(pluginlist)))
+    msg(("Installed plugins: {0}").format(", ".join(pluginlist)))
 
 
 def main():
diff --git a/src/livestreamer/compat.py b/src/livestreamer/compat.py
index 48a6718..e9c8c78 100644
--- a/src/livestreamer/compat.py
+++ b/src/livestreamer/compat.py
@@ -7,6 +7,7 @@ is_py3 = (sys.version_info[0] == 3)
 
 if is_py2:
     str = unicode
+    input = raw_input
 
     def bytes(b, enc="ascii"):
         return str(b)
@@ -14,6 +15,7 @@ if is_py2:
 elif is_py3:
     str = str
     bytes = bytes
+    input = input
 
 try:
     import urllib.request as urllib
diff --git a/src/livestreamer/plugins/justintv.py b/src/livestreamer/plugins/justintv.py
index 43e343c..c6d1219 100644
--- a/src/livestreamer/plugins/justintv.py
+++ b/src/livestreamer/plugins/justintv.py
@@ -110,7 +110,7 @@ class JustinTV(Plugin):
                 "swfUrl": self.SWFURL,
                 "swfhash": swfhash,
                 "swfsize": swfsize,
-                "live": 1
+                "live": True
             })
 
             if "token" in info:
diff --git a/src/livestreamer/plugins/ownedtv.py b/src/livestreamer/plugins/ownedtv.py
index 36887f7..d312efd 100644
--- a/src/livestreamer/plugins/ownedtv.py
+++ b/src/livestreamer/plugins/ownedtv.py
@@ -82,7 +82,7 @@ class OwnedTV(Plugin):
                 if not name in streams:
                     streams[name] = RTMPStream({
                         "rtmp": ("{0}/{1}").format(base, playpath),
-                        "live": 1
+                        "live": True
                     })
 
         return streams
diff --git a/src/livestreamer/plugins/ustreamtv.py b/src/livestreamer/plugins/ustreamtv.py
index 2e2e122..a7b88a7 100644
--- a/src/livestreamer/plugins/ustreamtv.py
+++ b/src/livestreamer/plugins/ustreamtv.py
@@ -46,7 +46,7 @@ class UStreamTV(Plugin):
                 "rtmp": ("{0}/{1}").format(cdnurl or fmsurl, playpath),
                 "pageUrl": self.url,
                 "swfUrl": self.SWFURL,
-                "live": 1
+                "live": True
             })
             streams["live"] = stream
 
diff --git a/src/livestreamer/stream.py b/src/livestreamer/stream.py
index c2b5658..235b01b 100644
--- a/src/livestreamer/stream.py
+++ b/src/livestreamer/stream.py
@@ -1,45 +1,35 @@
-from livestreamer.utils import CommandLine
+from livestreamer.utils import urlopen
 
-import subprocess, shlex
+import pbs
 
-class Stream(object):
-    def __init__(self, params={}):
-        self.params = params
-        self.process = None
+class StreamError(Exception):
+    pass
 
+class Stream(object):
     def open(self):
-        if self.process:
-            self.close()
-
-        cmdline = self.cmdline().format()
-        args = shlex.split(cmdline)
-
-        self.process = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+       raise NotImplementedError
 
-    def read(self, *args):
-        if self.process:
-            return self.process.stdout.read(*args)
+class RTMPStream(Stream):
+    def __init__(self, params):
+        self.params = params or {}
 
-    def close(self):
-        if self.process:
-            self.process.kill()
-            self.process = None
+    def open(self):
+        try:
+            rtmpdump = pbs.rtmpdump
+        except pbs.CommandNotFound:
+            raise StreamError("Unable to find 'rtmpdump' command")
 
-    def cmdline(self, out=None):
-       raise NotImplementedError
+        self.params["flv"] = "-"
+        self.params["_bg"] = True
 
-class RTMPStream(Stream):
-    def cmdline(self, out=None):
-        cmd = CommandLine("rtmpdump")
+        stream = rtmpdump(**self.params)
 
-        for key, value in self.params.items():
-            if key == "live":
-                if value == 1:
-                    cmd.args[key] = True
+        return stream.process.stdout
 
-            cmd.args[key] = value
+class HTTPStream(Stream):
+    def __init__(self, url):
+        self.url = url
 
-        if out:
-            cmd.args["flv"] = out
+    def open(self):
+        return urlopen(self.url)
 
-        return cmd
diff --git a/src/livestreamer/utils.py b/src/livestreamer/utils.py
index 5ab8a63..5735695 100644
--- a/src/livestreamer/utils.py
+++ b/src/livestreamer/utils.py
@@ -1,35 +1,11 @@
 #!/usr/bin/env python3
 
-from livestreamer.compat import urllib, bytes
+from livestreamer.compat import urllib
 from livestreamer.plugins import PluginError
 import hmac, hashlib, zlib, argparse
 
 SWF_KEY = b"Genuine Adobe Flash Player 001"
 
-class CommandLine(object):
-    def __init__(self, command):
-        self.command = command
-        self.args = {}
-        self.pipe = None
-
-    def format(self):
-        args = []
-
-        for key, value in self.args.items():
-            if value == True:
-                args.append(("--{0}").format(key))
-            else:
-                escaped = str(value).replace('"', '\\"').replace("$", "\$").replace("`", "\`")
-                args.append(("--{0} \"{1}\"").format(key, escaped))
-
-        args = (" ").join(args)
-        cmdline = ("{0} {1}").format(self.command, args)
-
-        if self.pipe:
-            cmdline += (" | {0}").format(self.pipe)
-
-        return cmdline
-
 class ArgumentParser(argparse.ArgumentParser):
     def convert_arg_line_to_args(self, line):
         split = line.find("=")
@@ -37,16 +13,27 @@ class ArgumentParser(argparse.ArgumentParser):
         val = line[split+1:].strip()
         yield "--%s=%s" % (key, val)
 
-def urlget(url, data=None, timeout=None, opener=None):
+def urlopen(url, data=None, timeout=None, opener=None):
     try:
         if opener is not None:
             fd = opener.open(url)
         else:
             fd = urllib.urlopen(url, data, timeout)
 
+    except IOError as err:
+        if type(err) is urllib.URLError:
+            raise PluginError(err.reason)
+        else:
+            raise PluginError(err)
+
+    return fd
+
+def urlget(url, data=None, timeout=None, opener=None):
+    fd = urlopen(url, data, timeout, opener)
+
+    try:
         data = fd.read()
         fd.close()
-
     except IOError as err:
         if type(err) is urllib.URLError:
             raise PluginError(err.reason)