now capture works for multiple arguments on pipe

extproc.py

@@ -394,7 +394,7 @@ def wait(self):
                     JOBS.remove(self)
 
 
-    def capture(self, fd):
+    def capture(self, *fd):
         """
         Fork-exec the Cmd and wait for its termination, capturing the
         output and/or error.
@@ -414,27 +414,32 @@ def capture(self, fd):
         Don't forget to close the file objects!
 
         """
+        assert len(fd) > 0
+        for descriptor in fd:
+            fd_update_dict = self._verify_capture_args(descriptor, self.fd)
+            self.fd.update(fd_update_dict)
+
+        if STDERR in fd:
+            self.fd[STDERR] = tempfile.TemporaryFile()
 
-        fd_update_dict = self._verify_capture_args(fd, self.fd)
-        self.fd.update(fd_update_dict)
         ## start piping
         prev = self.cmds[0].fd[0]
         for c in self.cmds[:-1]:
             if not _is_fileno(STDIN, c.fd[STDIN]):
                 prev = c.fd[STDIN]
-            if 2 == fd and _is_fileno(2, c.fd[2]):
-                c.fd[2] = self.fd[2]
+            if STDERR in fd and _is_fileno(STDERR, c.fd[STDERR]):
+                c.fd[STDERR] = self.fd[STDERR]
             c.p = c._popen(stdin=prev)
             prev = c.p.stdout
         ## prepare and fork the last child
         c = self.cmds[-1]
         if not _is_fileno(STDIN, c.fd[STDIN]):
             prev = c.fd[STDIN]
-        if STDOUT == fd:
+        if STDOUT in fd:
             ## we made sure that c.fd[1] had not been redirected before
             c.fd[STDOUT] = tempfile.TemporaryFile()
             self.fd[STDOUT] = c.fd[STDOUT]
-        if STDERR == fd and _is_fileno(STDERR, c.fd[STDERR]):
+        if STDERR in fd and _is_fileno(STDERR, c.fd[STDERR]):
             c.fd[STDERR] = self.fd[STDERR]
         c.p = c._popen(stdin=prev)
         ## wait for all children
@@ -444,9 +449,11 @@ def capture(self, fd):
         for c in self.cmds[:-1]:
             if c.fd[STDOUT] == PIPE:
               c.p.stdout.close()
-
-        self._cleanup_capture_dict(fd, self.fd)
-        self.fd[fd].seek(0)
+        if not set(fd) == set([1,2]):
+            #self._cleanup_capture(fd[0], p)
+            self._cleanup_capture_dict(fd[0], self.fd)
+        for descriptor in fd:
+            self.fd[descriptor].seek(0)
 
         return Capture(
             self.fd[1], self.fd[2], [c.p.returncode for c in self.cmds])

test_extproc.py

@@ -27,15 +27,6 @@ def test_sanity(self):
         self.assertEquals(run('true'), 0)
         self.assertEquals(run('false'), 1)
 
-    def test_Sh(self):
-        sh_call = Sh('echo bar >&2; echo foo; exit 1')
-        out, err, status = sh_call.capture(1, 2)
-        self.assertEquals(status, 1)
-        # these tests pass on OS X, not sure how they will run on
-        # linux
-        self.assertSh(out.read(), 'foo')
-        self.assertSh(err.read(), 'bar')
-
 
     def test_sh(self):
         """ test Cmd ENV """
@@ -45,23 +36,10 @@ def test_sh(self):
             NotImplementedError,
             lambda: sh('echo foo; echo bar >&2', {1: 2}))
 
-        ### test Cmd impossible capture
-        self.assertRaises(
-            ValueError,
-            lambda: sh("echo bogus stuff", {1: os.devnull}))
-
-        ### test Pipe stderr capture
-        pipe_ = Pipe(Sh('echo foo; sleep 0.01; echo  bar >&2'), Sh('cat >&2'))
-        self.assertSh(pipe_.capture(2).stderr.read(), 'foobar')
-
         ### test Pipe ENV
         self.assertSh(
             pipe(Sh('echo $x'), Sh('cat; echo $x'), e=dict(x='foobar')),
             'foobarfoobar')
-        ### test Pipe impossible capture
-        self.assertRaises(
-            ValueError,
-            lambda:pipe(Sh("echo bogus"), Cmd("cat", {1: os.devnull})))
 
         ### test Pipe pathetic case
         self.assertSh(pipe(Sh("echo foo"), Cmd("cat", {0: here("bar")})), 'bar')
@@ -79,7 +57,28 @@ def test_sh(self):
         f.seek(0)
         self.assertSh(f.read(), 'foo')
 
+    def test_capture(self):
+        sh_call = Sh('echo bar >&2; echo foo; exit 1')
+        out, err, status = sh_call.capture(1, 2)
+        self.assertEquals(status, 1)
+        # these tests pass on OS X, not sure how they will run on
+        # linux
+        self.assertSh(out.read(), 'foo')
+        self.assertSh(err.read(), 'bar')
+
+        ### test Pipe impossible capture
+        self.assertRaises(
+            ValueError,
+            lambda:pipe(Sh("echo bogus"), Cmd("cat", {1: os.devnull})))
+
+        ### test Cmd impossible capture
+        self.assertRaises(
+            ValueError,
+            lambda: sh("echo bogus stuff", {1: os.devnull}))
 
+        ### test Pipe stderr capture
+        pipe_ = Pipe(Sh('echo foo; sleep 0.01; echo  bar >&2'), Sh('cat >&2'))
+        self.assertSh(pipe_.capture(2).stderr.read(), 'foobar')
 
     def test_sh2(self):
         self.assertSh(sh('echo foo >&2', {STDERR: 1}), 'foo')
@@ -148,6 +147,14 @@ def test_capture(self):
                  Cmd('cat', {1: os.devnull})).capture(2).stderr.read(),
             'bar')
 
+        sh_call = Pipe(Sh('echo bar >&2; echo foo; exit 1'))
+        out, err, status = sh_call.capture(1, 2)
+        self.assertEquals(status, [1])
+        # these tests pass on OS X, not sure how they will run on
+        # linux
+        self.assertSh(out.read(), 'foo')
+        self.assertSh(err.read(), 'bar')
+
 class ExtProcCmdTest(ExtProcTest):
     def test_CMD(self):
         self.assertEquals(Cmd(['grep', 'my stuff']), Cmd('grep "my stuff"'))