Move the named pipe stuff to a file-like class.

2 files changed, 90 insertions, 68 deletions

diff --git a/src/livestreamer/cli.py b/src/livestreamer/cli.py
index ce54b37..d801b94 100644
--- a/src/livestreamer/cli.py
+++ b/src/livestreamer/cli.py
@@ -1,13 +1,13 @@
 import argparse
+import getpass
 import os
 import sys
 import subprocess
-import getpass
 
 from livestreamer import *
 from livestreamer.compat import input, stdout, is_win32
 from livestreamer.stream import StreamProcess
-from livestreamer.utils import ArgumentParser
+from livestreamer.utils import ArgumentParser, NamedPipe
 
 exampleusage = """
 example usage:
@@ -90,15 +90,6 @@ def set_msg_output(output):
     msg_output = output
     livestreamer.set_logoutput(output)
 
-def do_write(stream_out, stream_data):
-    if is_win32 and type(stream_out) is not file:
-        from ctypes import windll, cast, c_ulong, c_void_p, byref
-        windll.kernel32.ConnectNamedPipe(stream_out, None)
-        written = c_ulong(0)
-        windll.kernel32.WriteFile(stream_out, cast(stream_data, c_void_p), len(stream_data), byref(written), None)
-    else:
-        stream_out.write(stream_data)
-
 def write_stream(fd, out, progress, player):
     written = 0
 
@@ -119,7 +110,7 @@ def write_stream(fd, out, progress, player):
             break
 
         try:
-            do_write(out, data)
+            out.write(data)
         except IOError:
             logger.error("Error when writing to output")
             break
@@ -134,16 +125,9 @@ def write_stream(fd, out, progress, player):
 
     logger.info("Stream ended")
 
-    if is_win32 and type(out) is not file:
-        from ctypes import windll
-        windll.kernel32.DisconnectNamedPipe(out)
-    elif out != stdout:
-        try:
-            out.close()
-        except:
-            pass
-
 def check_output(output, force):
+    logger.debug("Checking output")
+
     if os.path.isfile(output) and not force:
         sys.stderr.write(("File {0} already exists! Overwrite it? [y/N] ").format(output))
 
@@ -169,20 +153,12 @@ def output_stream(stream, args):
     out = None
     player = None
 
-    if is_win32:
-        pipeprefix = "\\\\.\\pipe\\"
-    else:
-        import tempfile
-        pipeprefix = tempfile.gettempdir() + "/"
-
-    pipename = pipeprefix + "livestreamerpipe-" + str(os.getpid())
-
     logger.info("Opening stream: {0}", args.stream)
 
     try:
         fd = stream.open()
     except StreamError as err:
-        exit(("Could not open stream - {0}").format(err))
+        exit(("Could not open stream: {0}").format(err))
 
     logger.debug("Pre-buffering 8192 bytes")
     try:
@@ -193,8 +169,6 @@ def output_stream(stream, args):
     if len(prebuffer) == 0:
         exit("Failed to read data from stream")
 
-    logger.debug("Checking output")
-
     if args.output:
         if args.output == "-":
             out = stdout
@@ -205,33 +179,16 @@ def output_stream(stream, args):
         out = stdout
     else:
         if args.fifo:
+            pipename = "livestreamerpipe-" + str(os.getpid())
+
             logger.info("Creating pipe {0}", pipename)
 
-            cmd = args.player + " " + pipename
-            if is_win32:
-                from ctypes import windll
-                PIPE_ACCESS_OUTBOUND = 0x00000002
-                PIPE_TYPE_BYTE = 0x00000000
-                PIPE_READMODE_BYTE = 0x00000000
-                PIPE_WAIT = 0x00000000
-                PIPE_UNLIMITED_INSTANCES = 255
-                INVALID_HANDLE_VALUE = -1
-                bufsize = 8192
-
-                out = windll.kernel32.CreateNamedPipeA(pipename,
-                                                       PIPE_ACCESS_OUTBOUND,
-                                                       PIPE_TYPE_BYTE | PIPE_READMODE_BYTE | PIPE_WAIT,
-                                                       PIPE_UNLIMITED_INSTANCES,
-                                                       bufsize,
-                                                       bufsize,
-                                                       0,
-                                                       None)
-
-                if out == INVALID_HANDLE_VALUE:
-                    exit(("Failed to create pipe {0} - error code 0x{1:08X}").format(pipename, windll.kernel32.GetLastError()))
-            else:
-                os.mkfifo(pipename, 0o660)
+            try:
+                out = NamedPipe(pipename)
+            except IOError as err:
+                exit(("Failed to create pipe: {0}").format(err))
 
+            cmd = args.player + " " + out.path
             pin = sys.stdin
         else:
             cmd = args.player + " -"
@@ -249,11 +206,10 @@ def output_stream(stream, args):
                                   stdin=pin)
 
         if args.fifo:
-            if not is_win32:
-                try:
-                    out = open(pipename, "wb")
-                except IOError as err:
-                    exit(("Failed to open pipe {0} - {1}").format(pipename, err))
+            try:
+                out.open("wb")
+            except IOError as err:
+                exit(("Failed to open pipe {0} - {1}").format(pipename, err))
         else:
             out = player.stdin
 
@@ -265,7 +221,7 @@ def output_stream(stream, args):
         msvcrt.setmode(out.fileno(), os.O_BINARY)
 
     logger.debug("Writing stream to output")
-    do_write(out, prebuffer)
+    out.write(prebuffer)
 
     try:
         write_stream(fd, out, progress, player)
@@ -278,12 +234,12 @@ def output_stream(stream, args):
         except:
             pass
 
-    if args.fifo and not args.output and not args.stdout:
-        if is_win32:
-            from ctypes import windll
-            windll.kernel32.CloseHandle(out)
-        else:
-            os.unlink(pipename)
+    if out != stdout:
+        try:
+            out.close()
+        except:
+            pass
+
 
 def handle_url(args):
     try:
diff --git a/src/livestreamer/utils.py b/src/livestreamer/utils.py
index 015b7b2..a3633d9 100644
--- a/src/livestreamer/utils.py
+++ b/src/livestreamer/utils.py
@@ -1,11 +1,17 @@
+from .compat import is_win32
 from .plugins import PluginError
 
 import argparse
 import hashlib
 import hmac
+import os
 import requests
+import tempfile
 import zlib
 
+if is_win32:
+    from ctypes import windll, cast, c_ulong, c_void_p, byref
+
 SWFKey = b"Genuine Adobe Flash Player 001"
 RequestsConfig = { "danger_mode": True }
 
@@ -22,6 +28,66 @@ class ArgumentParser(argparse.ArgumentParser):
         else:
             yield "--%s" % line
 
+class NamedPipe(object):
+    def __init__(self, name):
+        self.fifo = None
+        self.pipe = None
+
+        if is_win32:
+            self.path = os.path.join("\\\\.\\pipe", name)
+            self.pipe = self._create_named_pipe(self.path)
+        else:
+            self.path = os.path.join(tempfile.gettempdir(), name)
+            self._create_fifo(self.path)
+
+    def _create_fifo(self, name):
+        os.mkfifo(name, 0o660)
+
+    def _create_named_pipe(self, path):
+        PIPE_ACCESS_OUTBOUND = 0x00000002
+        PIPE_TYPE_BYTE = 0x00000000
+        PIPE_READMODE_BYTE = 0x00000000
+        PIPE_WAIT = 0x00000000
+        PIPE_UNLIMITED_INSTANCES = 255
+        INVALID_HANDLE_VALUE = -1
+        bufsize = 8192
+
+        pipe = windll.kernel32.CreateNamedPipeA(path,
+                                                PIPE_ACCESS_OUTBOUND,
+                                                PIPE_TYPE_BYTE | PIPE_READMODE_BYTE | PIPE_WAIT,
+                                                PIPE_UNLIMITED_INSTANCES,
+                                                bufsize,
+                                                bufsize,
+                                                0,
+                                                None)
+
+        if pipe == INVALID_HANDLE_VALUE:
+            raise IOError(("error code 0x{0:08X}").format(windll.kernel32.GetLastError()))
+
+        return pipe
+
+    def open(self, mode):
+        if not self.pipe:
+            self.fifo = open(self.path, mode)
+
+    def write(self, data):
+        if self.pipe:
+            windll.kernel32.ConnectNamedPipe(self.pipe, None)
+            written = c_ulong(0)
+            windll.kernel32.WriteFile(self.pipe, cast(data, c_void_p),
+                                      len(data), byref(written),
+                                      None)
+            return written
+        else:
+            return self.fifo.write(data)
+
+    def close(self):
+        if self.pipe:
+            windll.kernel32.DisconnectNamedPipe(self.pipe)
+        else:
+            self.fifo.close()
+            os.unlink(self.path)
+
 def urlopen(url, method="get", exception=PluginError, **args):
     if "data" in args and args["data"] is not None:
         method = "post"