From dd39d7547c9fb60c1a6bec173e976c84061428c3 Mon Sep 17 00:00:00 2001
From: Jake VanderPlas <jakevdp@google.com>
Date: Fri, 27 Mar 2020 14:21:26 -0700
Subject: [PATCH] MAINT: add stderr_filter to _utils.check_output_with_stderr

---
 altair_saver/_utils.py           | 43 +++++++++++++++++++++++---------
 altair_saver/tests/test_utils.py | 31 +++++++++++++++++------
 2 files changed, 54 insertions(+), 20 deletions(-)

diff --git a/altair_saver/_utils.py b/altair_saver/_utils.py
index 68f5c16..3a0ed67 100644
--- a/altair_saver/_utils.py
+++ b/altair_saver/_utils.py
@@ -6,7 +6,7 @@
 import subprocess
 import sys
 import tempfile
-from typing import IO, Iterator, List, Optional, Union
+from typing import Callable, IO, Iterator, List, Optional, Union
 
 import altair as alt
 
@@ -168,31 +168,50 @@ def extract_format(fp: Union[IO, str]) -> str:
 
 
 def check_output_with_stderr(
-    cmd: Union[str, List[str]], shell: bool = False, input: Optional[bytes] = None
+    cmd: Union[str, List[str]], shell: bool = False, input: Optional[bytes] = None, stderr_filter:Callable[[str], bool] = None
 ) -> bytes:
     """Run a command in a subprocess, printing stderr to sys.stderr.
 
-    Arguments are passed directly to subprocess.run().
+    This function exists because normally, stderr from subprocess in the notebook
+    is printed to the terminal rather than to the notebook itself.
 
-    This is important because subprocess stderr in notebooks is printed to the
-    terminal rather than the notebook.
+    Parameters
+    ----------
+    cmd, shell, input :
+        Arguments are passed directly to `subprocess.run()`.
+    stderr_filter : function(str)->bool (optional)
+        If provided, this function is used to filter stderr lines from display.
+
+    Returns
+    -------
+    result : bytes
+        The stdout from the command
+
+    Raises
+    ------
+    subprocess.CalledProcessError : if the called process returns a non-zero exit code.
     """
     try:
         ps = subprocess.run(
             cmd,
             shell=shell,
+            input=input,
             check=True,
             stdout=subprocess.PIPE,
             stderr=subprocess.PIPE,
-            input=input,
         )
     except subprocess.CalledProcessError as err:
-        if err.stderr:
-            sys.stderr.write(err.stderr.decode())
-            sys.stderr.flush()
+        stderr = err.stderr
         raise
     else:
-        if ps.stderr:
-            sys.stderr.write(ps.stderr.decode())
-            sys.stderr.flush()
+        stderr = ps.stderr
         return ps.stdout
+    finally:
+        s = stderr.decode()
+        if stderr_filter:
+            s = '\n'.join(filter(stderr_filter, s.splitlines()))
+        if s:
+            if not s.endswith('\n'):
+                s += '\n'
+            sys.stderr.write(s)
+            sys.stderr.flush()
diff --git a/altair_saver/tests/test_utils.py b/altair_saver/tests/test_utils.py
index 6e2a357..f5613e5 100644
--- a/altair_saver/tests/test_utils.py
+++ b/altair_saver/tests/test_utils.py
@@ -5,7 +5,7 @@
 import tempfile
 
 import pytest
-from _pytest.capture import SysCaptureBinary
+from _pytest.capture import SysCapture
 
 from altair_saver.types import JSONDict
 from altair_saver._utils import (
@@ -142,10 +142,25 @@ def test_check_output_with_stderr(capsysbinary: SysCaptureBinary):
     assert captured.err == b"the error\n"
 
 
-def test_check_output_with_stderr_exit_1(capsysbinary: SysCaptureBinary):
-    with pytest.raises(subprocess.CalledProcessError) as err:
-        check_output_with_stderr(r'>&2 echo "the error" && exit 1', shell=True)
-    assert err.value.stderr == b"the error\n"
-    captured = capsysbinary.readouterr()
-    assert captured.out == b""
-    assert captured.err == b"the error\n"
+@pytest.mark.parametrize('cmd_error', [True, False])
+@pytest.mark.parametrize('use_filter', [True, False])
+def test_check_output_with_stderr(capsys: SysCapture, use_filter: bool, cmd_error: bool):
+    cmd = r'>&2 echo "first error\nsecond error" && echo "the output"'
+    stderr_filter = None if not use_filter else lambda line: line.startswith('second')
+
+    if cmd_error:
+        cmd += r' && exit 1'
+        with pytest.raises(subprocess.CalledProcessError) as err:
+            check_output_with_stderr(cmd, shell=True, stderr_filter=stderr_filter)
+        assert err.value.stderr == b"first error\nsecond error\n"
+    else:
+        output = check_output_with_stderr(cmd, shell=True, stderr_filter=stderr_filter)
+        assert output == b"the output\n"
+    
+    captured = capsys.readouterr()
+    assert captured.out == ""
+
+    if use_filter:
+        assert captured.err == "second error\n"
+    else:
+        assert captured.err == "first error\nsecond error\n"