[Process] Implement IteratorAggregate to stream output

nicolas-grekas · nicolas-grekas · commit cdd36278d0c8 · 2016-04-02T15:02:04.000+02:00
diff --git a/src/Symfony/Component/Process/Process.php b/src/Symfony/Component/Process/Process.php
@@ -27,7 +27,7 @@
  * @author Fabien Potencier <fabien@symfony.com>
  * @author Romain Neutron <imprec@gmail.com>
  */
-class Process
+class Process implements \IteratorAggregate
 {
     const ERR = 'err';
     const OUT = 'out';
@@ -362,7 +362,7 @@ public function wait(callable $callback = null)
         do {
             $this->checkTimeout();
             $running = '\\' === DIRECTORY_SEPARATOR ? $this->isRunning() : $this->processPipes->areOpen();
-            $this->readPipes($running, '\\' !== DIRECTORY_SEPARATOR || !$running);
+            $this->readPipes($running, '\\' !== DIRECTORY_SEPARATOR || !$running, $this->callback);
         } while ($running);
 
         while ($this->isRunning()) {
@@ -498,6 +498,54 @@ public function getIncrementalOutput()
         return $latest;
     }
 
+    /**
+     * Returns an iterator to the output of the process, with the output type as keys (Process::OUT/ERR).
+     *
+     * @param bool $blocking    Whether to use a blocking read call.
+     * @param bool $clearOutput Whether to clear or keep output in memory.
+     *
+     * @throws LogicException in case the output has been disabled
+     * @throws LogicException In case the process is not started
+     *
+     * @return \Generator
+     */
+    public function getIterator($blocking = true, $clearOutput = true)
+    {
+        $this->readPipesForOutput(__FUNCTION__, false);
+
+        while (null !== $this->callback) {
+            $out = stream_get_contents($this->stdout, -1, $this->incrementalOutputOffset);
+
+            if (isset($out[0])) {
+                if ($clearOutput) {
+                    $this->clearOutput();
+                } else {
+                    $this->incrementalOutputOffset = ftell($this->stdout);
+                }
+
+                yield self::OUT => $out;
+            } 
+
+            $err = stream_get_contents($this->stderr, -1, $this->incrementalErrorOutputOffset);
+
+            if (isset($err[0])) {
+                if ($clearOutput) {
+                    $this->clearErrorOutput();
+                } else {
+                    $this->incrementalErrorOutputOffset = ftell($this->stderr);
+                }
+
+                yield self::ERR => $err;
+            } 
+
+            if (!$blocking && !isset($out[0]) && !isset($err[0])) {
+                yield self::OUT => '';
+            }
+
+            $this->readPipesForOutput(__FUNCTION__, $blocking);
+        }
+    }
+
     /**
      * Clears the process output.
      *
@@ -1261,7 +1309,7 @@ protected function updateStatus($blocking)
         $this->processInformation = proc_get_status($this->process);
         $running = $this->processInformation['running'];
 
-        $this->readPipes($running && $blocking, '\\' !== DIRECTORY_SEPARATOR || !$running);
+        $this->readPipes($running && $blocking, '\\' !== DIRECTORY_SEPARATOR || !$running, $this->callback);
 
         if ($this->fallbackStatus && $this->enhanceSigchildCompatibility && $this->isSigchildEnabled()) {
             $this->processInformation = $this->fallbackStatus + $this->processInformation;
@@ -1296,19 +1344,20 @@ protected function isSigchildEnabled()
     /**
      * Reads pipes for the freshest output.
      *
-     * @param $caller The name of the method that needs fresh outputs
+     * @param string $caller   The name of the method that needs fresh outputs
+     * @param bool   $blocking Whether to use blocking calls or not.
      *
      * @throws LogicException in case output has been disabled or process is not started
      */
-    private function readPipesForOutput($caller)
+    private function readPipesForOutput($caller, $blocking = false)
     {
         if ($this->outputDisabled) {
             throw new LogicException('Output has been disabled.');
         }
 
         $this->requireProcessIsStarted($caller);
 
-        $this->updateStatus(false);
+        $this->updateStatus($blocking);
     }
 
     /**
@@ -1336,14 +1385,14 @@ private function validateTimeout($timeout)
     /**
      * Reads pipes, executes callback.
      *
-     * @param bool $blocking Whether to use blocking calls or not.
-     * @param bool $close    Whether to close file handles or not.
+     * @param bool     $blocking Whether to use blocking calls or not.
+     * @param bool     $close    Whether to close file handles or not.
+     * @param \Closure $callback The internal callback used to handle the data.
      */
-    private function readPipes($blocking, $close)
+    private function readPipes($blocking, $close, $callback)
     {
         $result = $this->processPipes->readAndWrite($blocking, $close);
 
-        $callback = $this->callback;
         foreach ($result as $type => $data) {
             if (3 !== $type) {
                 $callback($type === self::STDOUT ? self::OUT : self::ERR, $data);
diff --git a/src/Symfony/Component/Process/Tests/ProcessTest.php b/src/Symfony/Component/Process/Tests/ProcessTest.php
@@ -1270,6 +1270,36 @@ public function testInputStreamOnEmpty()
         $this->assertSame('123456', $process->getOutput());
     }
 
+    public function testIteratorOutput()
+    {
+        $process = new Process(self::$phpBin.' -r '.escapeshellarg('fwrite(STDOUT, 123); fwrite(STDERR, 234); flush(); usleep(100000); fwrite(STDOUT, 345); fwrite(STDERR, 456);'));
+        $process->start();
+        $output = array();
+
+        foreach ($process as $type => $data) {
+            $output[] = array($type, $data);
+            break;
+        }
+        $expectedOutput = array(
+            array($process::OUT, '123'),
+        );
+        $this->assertSame($expectedOutput, $output);
+
+        foreach ($process as $type => $data) {
+            $output[] = array($type, $data);
+        }
+
+        $this->assertFalse($process->isRunning());
+
+        $expectedOutput = array(
+            array($process::OUT, '123'),
+            array($process::ERR, '234'),
+            array($process::OUT, '345'),
+            array($process::ERR, '456'),
+        );
+        $this->assertSame($expectedOutput, $output);
+    }
+
     /**
      * @param string      $commandline
      * @param null|string $cwd
