Repository: tomerfiliba/plumbum

Issue:
TEE return code for pipeline is wrong when non-last command fails
Based on #145 
```py
from plumbum.cmd import yes, head
from plumbum import TEE, RETCODE

chain = yes['-x'] | head
print(chain & RETCODE)  # prints 1
with chain.bgrun(retcode=None) as p:
    print(p.run()[0])   # prints 1
print((chain & TEE(retcode=None))[0])  # prints 0
```

Comments:

Patch:
diff --git a/docs/index.rst b/docs/index.rst
index 78bf283..c1b2ed0 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -166,7 +166,7 @@ that I wrote for `RPyC <https://rpyc.readthedocs.io/>`_. When I combined the two
 
 Credits
 =======
-The project has been inspired by **PBS** (now called `sh <http://amoffat.github.io/sh/>`_)
+The project has been inspired by **PBS** (now called `sh <http://sh.rtfd.org>`_)
 of `Andrew Moffat <https://github.com/amoffat>`_,
 and has borrowed some of his ideas (namely treating programs like functions and the
 nice trick for importing commands). However, I felt there was too much magic going on in PBS,
diff --git a/plumbum/commands/base.py b/plumbum/commands/base.py
index 52c0f26..3172c52 100644
--- a/plumbum/commands/base.py
+++ b/plumbum/commands/base.py
@@ -398,8 +398,6 @@ class Pipeline(BaseCommand):
         dstproc = self.dstcmd.popen(**kwargs)
         # allow p1 to receive a SIGPIPE if p2 exits
         srcproc.stdout.close()
-        if srcproc.stderr is not None:
-            dstproc.stderr = srcproc.stderr
         if srcproc.stdin and src_kwargs.get("stdin") != PIPE:
             srcproc.stdin.close()
         dstproc.srcproc = srcproc
diff --git a/plumbum/commands/modifiers.py b/plumbum/commands/modifiers.py
index 98b3749..af59071 100644
--- a/plumbum/commands/modifiers.py
+++ b/plumbum/commands/modifiers.py
@@ -226,6 +226,7 @@ class _TEE(ExecutionModifier):
 
                         buf.append(data)
 
+            p.wait()  # To get return code in p
             stdout = "".join([x.decode("utf-8") for x in outbuf])
             stderr = "".join([x.decode("utf-8") for x in errbuf])
             return p.returncode, stdout, stderr
diff --git a/plumbum/commands/processes.py b/plumbum/commands/processes.py
index 802ede4..89c2763 100644
--- a/plumbum/commands/processes.py
+++ b/plumbum/commands/processes.py
@@ -18,14 +18,44 @@ def _check_process(proc, retcode, timeout, stdout, stderr):
     return proc.returncode, stdout, stderr
 
 
+def _get_piped_streams(proc):
+    """Get a list of all valid standard streams for proc that were opened with PIPE option.
+
+    If proc was started from a Pipeline command, this function assumes it will have a
+    "srcproc" member pointing to the previous command in the pipeline. That link will
+    be used to traverse all started processes started from the pipeline, the list will
+    include stdout/stderr streams opened as PIPE for all commands in the pipeline.
+    If that was not the case, some processes could write to pipes no one reads from
+    which would result in process stalling after the pipe's buffer is filled.
+
+    Streams that were closed (because they were redirected to the input of a subsequent command)
+    are not included in the result
+    """
+    streams = []
+
+    def add_stream(type_, stream):
+        if stream is None or stream.closed:
+            return
+        streams.append((type_, stream))
+
+    while proc:
+        add_stream(1, proc.stderr)
+        add_stream(0, proc.stdout)
+        proc = getattr(proc, "srcproc", None)
+
+    return streams
+
+
 def _iter_lines_posix(proc, decode, linesize, line_timeout=None):
     from selectors import EVENT_READ, DefaultSelector
 
+    streams = _get_piped_streams(proc)
+
     # Python 3.4+ implementation
     def selector():
         sel = DefaultSelector()
-        sel.register(proc.stdout, EVENT_READ, 0)
-        sel.register(proc.stderr, EVENT_READ, 1)
+        for stream_type, stream in streams:
+            sel.register(stream, EVENT_READ, stream_type)
         while True:
             ready = sel.select(line_timeout)
             if not ready and line_timeout:
@@ -41,10 +71,9 @@ def _iter_lines_posix(proc, decode, linesize, line_timeout=None):
         yield ret
         if proc.poll() is not None:
             break
-    for line in proc.stdout:
-        yield 0, decode(line)
-    for line in proc.stderr:
-        yield 1, decode(line)
+    for stream_type, stream in streams:
+        for line in stream:
+            yield stream_type, decode(line)
 
 
 def _iter_lines_win32(proc, decode, linesize, line_timeout=None):


Test patch:
diff --git a/tests/test_local.py b/tests/test_local.py
index 889624a..3b5eb51 100644
--- a/tests/test_local.py
+++ b/tests/test_local.py
@@ -607,7 +607,7 @@ class TestLocalMachine:
 
     @skip_on_windows
     def test_modifiers(self):
-        from plumbum.cmd import grep, ls
+        from plumbum.cmd import cat, grep, ls
 
         f = (ls["-a"] | grep["\\.py"]) & BG
         f.wait()
@@ -615,11 +615,17 @@ class TestLocalMachine:
 
         command = ls["-a"] | grep["local"]
         command_false = ls["-a"] | grep["not_a_file_here"]
+        command_false_2 = command_false | cat
         command & FG
         assert command & TF
         assert not (command_false & TF)
+        assert not (command_false_2 & TF)
         assert command & RETCODE == 0
         assert command_false & RETCODE == 1
+        assert command_false_2 & RETCODE == 1
+        assert (command & TEE)[0] == 0
+        assert (command_false & TEE(retcode=None))[0] == 1
+        assert (command_false_2 & TEE(retcode=None))[0] == 1
 
     @skip_on_windows
     def test_tee_modifier(self, capfd):
diff --git a/tests/test_pipelines.py b/tests/test_pipelines.py
new file mode 100644
index 0000000..8ea47e2
--- /dev/null
+++ b/tests/test_pipelines.py
@@ -0,0 +1,87 @@
+from typing import List, Tuple
+
+import pytest
+
+import plumbum
+from plumbum._testtools import skip_on_windows
+from plumbum.commands import BaseCommand
+
+
+@skip_on_windows
+@pytest.mark.timeout(3)
+def test_draining_stderr(generate_cmd, process_cmd):
+    stdout, stderr = get_output_with_iter_lines(
+        generate_cmd | process_cmd | process_cmd
+    )
+    expected_output = {f"generated {i}" for i in range(5000)}
+    expected_output.update(f"consumed {i}" for i in range(5000))
+    assert set(stderr) - expected_output == set()
+    assert len(stderr) == 15000
+    assert len(stdout) == 5000
+
+
+@skip_on_windows
+@pytest.mark.timeout(3)
+def test_draining_stderr_with_stderr_redirect(tmp_path, generate_cmd, process_cmd):
+    stdout, stderr = get_output_with_iter_lines(
+        generate_cmd | (process_cmd >= str(tmp_path / "output.txt")) | process_cmd
+    )
+    expected_output = {f"generated {i}" for i in range(5000)}
+    expected_output.update(f"consumed {i}" for i in range(5000))
+    assert set(stderr) - expected_output == set()
+    assert len(stderr) == 10000
+    assert len(stdout) == 5000
+
+
+@skip_on_windows
+@pytest.mark.timeout(3)
+def test_draining_stderr_with_stdout_redirect(tmp_path, generate_cmd, process_cmd):
+    stdout, stderr = get_output_with_iter_lines(
+        generate_cmd | process_cmd | process_cmd > str(tmp_path / "output.txt")
+    )
+    expected_output = {f"generated {i}" for i in range(5000)}
+    expected_output.update(f"consumed {i}" for i in range(5000))
+    assert set(stderr) - expected_output == set()
+    assert len(stderr) == 15000
+    assert len(stdout) == 0
+
+
+@pytest.fixture()
+def generate_cmd(tmp_path):
+    generate = tmp_path / "generate.py"
+    generate.write_text(
+        """\
+import sys
+for i in range(5000):
+    print("generated", i, file=sys.stderr)
+    print(i)
+"""
+    )
+    return plumbum.local["python"][generate]
+
+
+@pytest.fixture()
+def process_cmd(tmp_path):
+    process = tmp_path / "process.py"
+    process.write_text(
+        """\
+import sys
+for line in sys.stdin:
+    i = line.strip()
+    print("consumed", i, file=sys.stderr)
+    print(i)
+"""
+    )
+    return plumbum.local["python"][process]
+
+
+def get_output_with_iter_lines(cmd: BaseCommand) -> Tuple[List[str], List[str]]:
+    stderr, stdout = [], []
+    proc = cmd.popen()
+    for stdout_line, stderr_line in proc.iter_lines(retcode=[0, None]):
+        if stderr_line:
+            stderr.append(stderr_line)
+        if stdout_line:
+            stdout.append(stdout_line)
+    proc.wait()
+    return stdout, stderr