Utility for redirecting C++ streams to Python (#1009)

diff --git a/docs/advanced/pycpp/utilities.rst b/docs/advanced/pycpp/utilities.rst
index 171a885..369e7c9 100644
--- a/docs/advanced/pycpp/utilities.rst
+++ b/docs/advanced/pycpp/utilities.rst
@@ -21,6 +21,72 @@
     auto args = py::make_tuple("unpacked", true);
     py::print("->", *args, "end"_a="<-"); // -> unpacked True <-
 
+.. _ostream_redirect:
+
+Capturing standard output from ostream
+======================================
+
+Often, a library will use the streams ``std::cout`` and ``std::cerr`` to print,
+but this does not play well with Python's standard ``sys.stdout`` and ``sys.stderr``
+redirection. Replacing a library's printing with `py::print <print>` may not
+be feasible. This can be fixed using a guard around the library function that
+redirects output to the corresponding Python streams:
+
+.. code-block:: cpp
+
+    #include <pybind11/iostream.h>
+
+    ...
+
+    // Add a scoped redirect for your noisy code
+    m.def("noisy_func", []() {
+        py::scoped_ostream_redirect stream(
+            std::cout,                               // std::ostream&
+            py::module::import("sys").attr("stdout") // Python output
+        );
+        call_noisy_func();
+    });
+
+This method respects flushes on the output streams and will flush if needed
+when the scoped guard is destroyed. This allows the output to be redirected in
+real time, such as to a Jupyter notebook. The two arguments, the C++ stream and
+the Python output, are optional, and default to standard output if not given. An
+extra type, `py::scoped_estream_redirect <scoped_estream_redirect>`, is identical
+except for defaulting to ``std::cerr`` and ``sys.stderr``; this can be useful with
+`py::call_guard`, which allows multiple items, but uses the default constructor:
+
+.. code-block:: py
+
+    // Alternative: Call single function using call guard
+    m.def("noisy_func", &call_noisy_function,
+          py::call_guard<py::scoped_ostream_redirect,
+                         py::scoped_estream_redirect>());
+
+The redirection can also be done in Python with the addition of a context
+manager, using the `py::add_ostream_redirect() <add_ostream_redirect>` function:
+
+.. code-block:: cpp
+
+    py::add_ostream_redirect(m, "ostream_redirect");
+
+The name in Python defaults to ``ostream_redirect`` if no name is passed.  This
+creates the following context manager in Python:
+
+.. code-block:: python
+
+    with ostream_redirect(stdout=True, stderr=True):
+        noisy_function()
+
+It defaults to redirecting both streams, though you can use the keyword
+arguments to disable one of the streams if needed.
+
+.. note::
+
+    The above methods will not redirect C-level output to file descriptors, such
+    as ``fprintf``. For those cases, you'll need to redirect the file
+    descriptors either directly in C or with Python's ``os.dup2`` function
+    in an operating-system dependent way.
+
 .. _eval:
 
 Evaluating Python expressions from strings and files
diff --git a/docs/changelog.rst b/docs/changelog.rst
index 2e585dc..49249e3 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -123,10 +123,15 @@
   7. Fixed lifetime of temporary C++ objects created in Python-to-C++ conversions.
      `#924 <https://github.com/pybind/pybind11/pull/924>`_.
 
-* Scope guard call policy for RAII types, e.g. ``py::call_guard<py::gil_scoped_release>()``.
-  See :ref:`call_policies` for details.
+* Scope guard call policy for RAII types, e.g. ``py::call_guard<py::gil_scoped_release>()``,
+  ``py::call_guard<py::scoped_ostream_redirect>()``. See :ref:`call_policies` for details.
   `#740 <https://github.com/pybind/pybind11/pull/740>`_.
 
+* Utility for redirecting C++ streams to Python (e.g. ``std::cout`` ->
+  ``sys.stdout``). Scope guard ``py::scoped_ostream_redirect`` in C++ and
+  a context manager in Python. See :ref:`ostream_redirect`.
+  `#1009 <https://github.com/pybind/pybind11/pull/1009>`_.
+
 * Improved handling of types and exceptions across module boundaries.
   `#915 <https://github.com/pybind/pybind11/pull/915>`_,
   `#951 <https://github.com/pybind/pybind11/pull/951>`_,
@@ -298,7 +303,6 @@
   `#923 <https://github.com/pybind/pybind11/pull/923>`_,
   `#963 <https://github.com/pybind/pybind11/pull/963>`_.
 
-
 v2.1.1 (April 7, 2017)
 -----------------------------------------------------
 
diff --git a/docs/reference.rst b/docs/reference.rst
index 1365739..e41141b 100644
--- a/docs/reference.rst
+++ b/docs/reference.rst
@@ -71,6 +71,15 @@
 
 .. doxygenclass:: scoped_interpreter
 
+Redirecting C++ streams
+=======================
+
+.. doxygenclass:: scoped_ostream_redirect
+
+.. doxygenclass:: scoped_estream_redirect
+
+.. doxygenfunction:: add_ostream_redirect
+
 Python build-in functions
 =========================
 
diff --git a/include/pybind11/iostream.h b/include/pybind11/iostream.h
new file mode 100644
index 0000000..a9c27aa
--- /dev/null
+++ b/include/pybind11/iostream.h
@@ -0,0 +1,200 @@
+/*
+    pybind11/iostream.h -- Tools to assist with redirecting cout and cerr to Python
+
+    Copyright (c) 2017 Henry F. Schreiner
+
+    All rights reserved. Use of this source code is governed by a
+    BSD-style license that can be found in the LICENSE file.
+*/
+
+#pragma once
+
+#include "pybind11.h"
+
+#include <streambuf>
+#include <ostream>
+#include <string>
+#include <memory>
+#include <iostream>
+
+NAMESPACE_BEGIN(PYBIND11_NAMESPACE)
+NAMESPACE_BEGIN(detail)
+
+// Buffer that writes to Python instead of C++
+class pythonbuf : public std::streambuf {
+private:
+    using traits_type = std::streambuf::traits_type;
+
+    char d_buffer[1024];
+    object pywrite;
+    object pyflush;
+
+    int overflow(int c) {
+        if (!traits_type::eq_int_type(c, traits_type::eof())) {
+            *pptr() = traits_type::to_char_type(c);
+            pbump(1);
+        }
+        return sync() ? traits_type::not_eof(c) : traits_type::eof();
+    }
+
+    int sync() {
+        if (pbase() != pptr()) {
+            // This subtraction cannot be negative, so dropping the sign
+            str line(pbase(), static_cast<size_t>(pptr() - pbase()));
+
+            pywrite(line);
+            pyflush();
+
+            setp(pbase(), epptr());
+        }
+        return 0;
+    }
+
+public:
+    pythonbuf(object pyostream)
+        : pywrite(pyostream.attr("write")),
+          pyflush(pyostream.attr("flush")) {
+        setp(d_buffer, d_buffer + sizeof(d_buffer) - 1);
+    }
+
+    /// Sync before destroy
+    ~pythonbuf() {
+        sync();
+    }
+};
+
+NAMESPACE_END(detail)
+
+
+/** \rst
+    This a move-only guard that redirects output.
+
+    .. code-block:: cpp
+
+        #include <pybind11/iostream.h>
+
+        ...
+
+        {
+            py::scoped_ostream_redirect output;
+            std::cout << "Hello, World!"; // Python stdout
+        } // <-- return std::cout to normal
+
+    You can explicitly pass the c++ stream and the python object,
+    for example to guard stderr instead.
+
+    .. code-block:: cpp
+
+        {
+            py::scoped_ostream_redirect output{std::cerr, py::module::import("sys").attr("stderr")};
+            std::cerr << "Hello, World!";
+        }
+ \endrst */
+class scoped_ostream_redirect {
+protected:
+    std::streambuf *old;
+    std::ostream &costream;
+    detail::pythonbuf buffer;
+
+public:
+    scoped_ostream_redirect(
+            std::ostream &costream = std::cout,
+            object pyostream = module::import("sys").attr("stdout"))
+        : costream(costream), buffer(pyostream) {
+        old = costream.rdbuf(&buffer);
+    }
+
+    ~scoped_ostream_redirect() {
+        costream.rdbuf(old);
+    }
+
+    scoped_ostream_redirect(const scoped_ostream_redirect &) = delete;
+    scoped_ostream_redirect(scoped_ostream_redirect &&other) = default;
+    scoped_ostream_redirect &operator=(const scoped_ostream_redirect &) = delete;
+    scoped_ostream_redirect &operator=(scoped_ostream_redirect &&) = delete;
+};
+
+
+/** \rst
+    Like `scoped_ostream_redirect`, but redirects cerr by default. This class
+    is provided primary to make ``py::call_guard`` easier to make.
+
+    .. code-block:: cpp
+
+     m.def("noisy_func", &noisy_func,
+           py::call_guard<scoped_ostream_redirect,
+                          scoped_estream_redirect>());
+
+\endrst */
+class scoped_estream_redirect : public scoped_ostream_redirect {
+public:
+    scoped_estream_redirect(
+            std::ostream &costream = std::cerr,
+            object pyostream = module::import("sys").attr("stderr"))
+        : scoped_ostream_redirect(costream,pyostream) {}
+};
+
+
+NAMESPACE_BEGIN(detail)
+
+// Class to redirect output as a context manager. C++ backend.
+class OstreamRedirect {
+    bool do_stdout_;
+    bool do_stderr_;
+    std::unique_ptr<scoped_ostream_redirect> redirect_stdout;
+    std::unique_ptr<scoped_estream_redirect> redirect_stderr;
+
+public:
+    OstreamRedirect(bool do_stdout = true, bool do_stderr = true)
+        : do_stdout_(do_stdout), do_stderr_(do_stderr) {}
+
+    void enter() {
+        if (do_stdout_)
+            redirect_stdout.reset(new scoped_ostream_redirect());
+        if (do_stderr_)
+            redirect_stderr.reset(new scoped_estream_redirect());
+    }
+
+    void exit() {
+        redirect_stdout.reset();
+        redirect_stderr.reset();
+    }
+};
+
+NAMESPACE_END(detail)
+
+/** \rst
+    This is a helper function to add a C++ redirect context manager to Python
+    instead of using a C++ guard. To use it, add the following to your binding code:
+
+    .. code-block:: cpp
+
+        #include <pybind11/iostream.h>
+
+        ...
+
+        py::add_ostream_redirect(m, "ostream_redirect");
+
+    You now have a Python context manager that redirects your output:
+
+    .. code-block:: python
+
+        with m.ostream_redirect():
+            m.print_to_cout_function()
+
+    This manager can optionally be told which streams to operate on:
+
+    .. code-block:: python
+
+        with m.ostream_redirect(stdout=true, stderr=true):
+            m.noisy_function_with_error_printing()
+
+ \endrst */
+inline class_<detail::OstreamRedirect> add_ostream_redirect(module m, std::string name = "ostream_redirect") {
+    return class_<detail::OstreamRedirect>(m, name.c_str(), module_local())
+        .def(init<bool,bool>(), arg("stdout")=true, arg("stderr")=true)
+        .def("__enter__", &detail::OstreamRedirect::enter)
+        .def("__exit__", [](detail::OstreamRedirect &self, args) { self.exit(); });
+}
+
+NAMESPACE_END(PYBIND11_NAMESPACE)
diff --git a/setup.py b/setup.py
index 6e7a40a..e24f34b 100644
--- a/setup.py
+++ b/setup.py
@@ -28,6 +28,7 @@
         'include/pybind11/embed.h',
         'include/pybind11/eval.h',
         'include/pybind11/functional.h',
+        'include/pybind11/iostream.h',
         'include/pybind11/numpy.h',
         'include/pybind11/operators.h',
         'include/pybind11/options.h',
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index 1e68a74..0d6ebfc 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -40,6 +40,7 @@
   test_eval.cpp
   test_exceptions.cpp
   test_factory_constructors.cpp
+  test_iostream.cpp
   test_kwargs_and_defaults.cpp
   test_local_bindings.cpp
   test_methods_and_attributes.cpp
diff --git a/tests/test_iostream.cpp b/tests/test_iostream.cpp
new file mode 100644
index 0000000..e67f88a
--- /dev/null
+++ b/tests/test_iostream.cpp
@@ -0,0 +1,73 @@
+/*
+    tests/test_iostream.cpp -- Usage of scoped_output_redirect
+
+    Copyright (c) 2017 Henry F. Schreiner
+
+    All rights reserved. Use of this source code is governed by a
+    BSD-style license that can be found in the LICENSE file.
+*/
+
+
+#include <pybind11/iostream.h>
+#include "pybind11_tests.h"
+#include <iostream>
+
+
+void noisy_function(std::string msg, bool flush) {
+
+    std::cout << msg;
+    if (flush)
+        std::cout << std::flush;
+}
+
+void noisy_funct_dual(std::string msg, std::string emsg) {
+    std::cout << msg;
+    std::cerr << emsg;
+}
+
+TEST_SUBMODULE(iostream, m) {
+
+    add_ostream_redirect(m);
+
+    // test_evals
+
+    m.def("captured_output_default", [](std::string msg) {
+        py::scoped_ostream_redirect redir;
+        std::cout << msg << std::flush;
+    });
+
+    m.def("captured_output", [](std::string msg) {
+        py::scoped_ostream_redirect redir(std::cout, py::module::import("sys").attr("stdout"));
+        std::cout << msg << std::flush;
+    });
+
+    m.def("guard_output", &noisy_function,
+            py::call_guard<py::scoped_ostream_redirect>(),
+            py::arg("msg"), py::arg("flush")=true);
+
+    m.def("captured_err", [](std::string msg) {
+        py::scoped_ostream_redirect redir(std::cerr, py::module::import("sys").attr("stderr"));
+        std::cerr << msg << std::flush;
+    });
+
+    m.def("noisy_function", &noisy_function, py::arg("msg"), py::arg("flush") = true);
+
+    m.def("dual_guard", &noisy_funct_dual,
+            py::call_guard<py::scoped_ostream_redirect, py::scoped_estream_redirect>(),
+            py::arg("msg"), py::arg("emsg"));
+
+    m.def("raw_output", [](std::string msg) {
+        std::cout << msg << std::flush;
+    });
+
+    m.def("raw_err", [](std::string msg) {
+        std::cerr << msg << std::flush;
+    });
+
+    m.def("captured_dual", [](std::string msg, std::string emsg) {
+        py::scoped_ostream_redirect redirout(std::cout, py::module::import("sys").attr("stdout"));
+        py::scoped_ostream_redirect redirerr(std::cerr, py::module::import("sys").attr("stderr"));
+        std::cout << msg << std::flush;
+        std::cerr << emsg << std::flush;
+    });
+}
diff --git a/tests/test_iostream.py b/tests/test_iostream.py
new file mode 100644
index 0000000..3364849
--- /dev/null
+++ b/tests/test_iostream.py
@@ -0,0 +1,203 @@
+from pybind11_tests import iostream as m
+import sys
+
+from contextlib import contextmanager
+
+try:
+    # Python 3
+    from io import StringIO
+except ImportError:
+    # Python 2
+    try:
+        from cStringIO import StringIO
+    except ImportError:
+        from StringIO import StringIO
+
+try:
+    # Python 3.4
+    from contextlib import redirect_stdout
+except ImportError:
+    @contextmanager
+    def redirect_stdout(target):
+        original = sys.stdout
+        sys.stdout = target
+        yield
+        sys.stdout = original
+
+try:
+    # Python 3.5
+    from contextlib import redirect_stderr
+except ImportError:
+    @contextmanager
+    def redirect_stderr(target):
+        original = sys.stderr
+        sys.stderr = target
+        yield
+        sys.stderr = original
+
+
+def test_captured(capsys):
+    msg = "I've been redirected to Python, I hope!"
+    m.captured_output(msg)
+    stdout, stderr = capsys.readouterr()
+    assert stdout == msg
+    assert stderr == ''
+
+    m.captured_output_default(msg)
+    stdout, stderr = capsys.readouterr()
+    assert stdout == msg
+    assert stderr == ''
+
+    m.captured_err(msg)
+    stdout, stderr = capsys.readouterr()
+    assert stdout == ''
+    assert stderr == msg
+
+
+def test_guard_capture(capsys):
+    msg = "I've been redirected to Python, I hope!"
+    m.guard_output(msg)
+    stdout, stderr = capsys.readouterr()
+    assert stdout == msg
+    assert stderr == ''
+
+
+def test_series_captured(capture):
+    with capture:
+        m.captured_output("a")
+        m.captured_output("b")
+    assert capture == "ab"
+
+
+def test_flush(capfd):
+    msg = "(not flushed)"
+    msg2 = "(flushed)"
+
+    with m.ostream_redirect():
+        m.noisy_function(msg, flush=False)
+        stdout, stderr = capfd.readouterr()
+        assert stdout == ''
+
+        m.noisy_function(msg2, flush=True)
+        stdout, stderr = capfd.readouterr()
+        assert stdout == msg + msg2
+
+        m.noisy_function(msg, flush=False)
+
+    stdout, stderr = capfd.readouterr()
+    assert stdout == msg
+
+
+def test_not_captured(capfd):
+    msg = "Something that should not show up in log"
+    stream = StringIO()
+    with redirect_stdout(stream):
+        m.raw_output(msg)
+    stdout, stderr = capfd.readouterr()
+    assert stdout == msg
+    assert stderr == ''
+    assert stream.getvalue() == ''
+
+    stream = StringIO()
+    with redirect_stdout(stream):
+        m.captured_output(msg)
+    stdout, stderr = capfd.readouterr()
+    assert stdout == ''
+    assert stderr == ''
+    assert stream.getvalue() == msg
+
+
+def test_err(capfd):
+    msg = "Something that should not show up in log"
+    stream = StringIO()
+    with redirect_stderr(stream):
+        m.raw_err(msg)
+    stdout, stderr = capfd.readouterr()
+    assert stdout == ''
+    assert stderr == msg
+    assert stream.getvalue() == ''
+
+    stream = StringIO()
+    with redirect_stderr(stream):
+        m.captured_err(msg)
+    stdout, stderr = capfd.readouterr()
+    assert stdout == ''
+    assert stderr == ''
+    assert stream.getvalue() == msg
+
+
+def test_multi_captured(capfd):
+    stream = StringIO()
+    with redirect_stdout(stream):
+        m.captured_output("a")
+        m.raw_output("b")
+        m.captured_output("c")
+        m.raw_output("d")
+    stdout, stderr = capfd.readouterr()
+    assert stdout == 'bd'
+    assert stream.getvalue() == 'ac'
+
+
+def test_dual(capsys):
+    m.captured_dual("a", "b")
+    stdout, stderr = capsys.readouterr()
+    assert stdout == "a"
+    assert stderr == "b"
+
+
+def test_redirect(capfd):
+    msg = "Should not be in log!"
+    stream = StringIO()
+    with redirect_stdout(stream):
+        m.raw_output(msg)
+    stdout, stderr = capfd.readouterr()
+    assert stdout == msg
+    assert stream.getvalue() == ''
+
+    stream = StringIO()
+    with redirect_stdout(stream):
+        with m.ostream_redirect():
+            m.raw_output(msg)
+    stdout, stderr = capfd.readouterr()
+    assert stdout == ''
+    assert stream.getvalue() == msg
+
+    stream = StringIO()
+    with redirect_stdout(stream):
+        m.raw_output(msg)
+    stdout, stderr = capfd.readouterr()
+    assert stdout == msg
+    assert stream.getvalue() == ''
+
+
+def test_redirect_err(capfd):
+    msg = "StdOut"
+    msg2 = "StdErr"
+
+    stream = StringIO()
+    with redirect_stderr(stream):
+        with m.ostream_redirect(stdout=False):
+            m.raw_output(msg)
+            m.raw_err(msg2)
+    stdout, stderr = capfd.readouterr()
+    assert stdout == msg
+    assert stderr == ''
+    assert stream.getvalue() == msg2
+
+
+def test_redirect_both(capfd):
+    msg = "StdOut"
+    msg2 = "StdErr"
+
+    stream = StringIO()
+    stream2 = StringIO()
+    with redirect_stdout(stream):
+        with redirect_stderr(stream2):
+            with m.ostream_redirect():
+                m.raw_output(msg)
+                m.raw_err(msg2)
+    stdout, stderr = capfd.readouterr()
+    assert stdout == ''
+    assert stderr == ''
+    assert stream.getvalue() == msg
+    assert stream2.getvalue() == msg2