iostream.h revision 14299:2fbea9df56d2 
1/*
2    pybind11/iostream.h -- Tools to assist with redirecting cout and cerr to Python
3
4    Copyright (c) 2017 Henry F. Schreiner
5
6    All rights reserved. Use of this source code is governed by a
7    BSD-style license that can be found in the LICENSE file.
8*/
9
10#pragma once
11
12#include "pybind11.h"
13
14#include <streambuf>
15#include <ostream>
16#include <string>
17#include <memory>
18#include <iostream>
19
20NAMESPACE_BEGIN(PYBIND11_NAMESPACE)
21NAMESPACE_BEGIN(detail)
22
23// Buffer that writes to Python instead of C++
24class pythonbuf : public std::streambuf {
25private:
26    using traits_type = std::streambuf::traits_type;
27
28    const size_t buf_size;
29    std::unique_ptr<char[]> d_buffer;
30    object pywrite;
31    object pyflush;
32
33    int overflow(int c) {
34        if (!traits_type::eq_int_type(c, traits_type::eof())) {
35            *pptr() = traits_type::to_char_type(c);
36            pbump(1);
37        }
38        return sync() == 0 ? traits_type::not_eof(c) : traits_type::eof();
39    }
40
41    int sync() {
42        if (pbase() != pptr()) {
43            // This subtraction cannot be negative, so dropping the sign
44            str line(pbase(), static_cast<size_t>(pptr() - pbase()));
45
46            {
47                gil_scoped_acquire tmp;
48                pywrite(line);
49                pyflush();
50            }
51
52            setp(pbase(), epptr());
53        }
54        return 0;
55    }
56
57public:
58
59    pythonbuf(object pyostream, size_t buffer_size = 1024)
60        : buf_size(buffer_size),
61          d_buffer(new char[buf_size]),
62          pywrite(pyostream.attr("write")),
63          pyflush(pyostream.attr("flush")) {
64        setp(d_buffer.get(), d_buffer.get() + buf_size - 1);
65    }
66
67    pythonbuf(pythonbuf&&) = default;
68
69    /// Sync before destroy
70    ~pythonbuf() {
71        sync();
72    }
73};
74
75NAMESPACE_END(detail)
76
77
78/** \rst
79    This a move-only guard that redirects output.
80
81    .. code-block:: cpp
82
83        #include <pybind11/iostream.h>
84
85        ...
86
87        {
88            py::scoped_ostream_redirect output;
89            std::cout << "Hello, World!"; // Python stdout
90        } // <-- return std::cout to normal
91
92    You can explicitly pass the c++ stream and the python object,
93    for example to guard stderr instead.
94
95    .. code-block:: cpp
96
97        {
98            py::scoped_ostream_redirect output{std::cerr, py::module::import("sys").attr("stderr")};
99            std::cerr << "Hello, World!";
100        }
101 \endrst */
102class scoped_ostream_redirect {
103protected:
104    std::streambuf *old;
105    std::ostream &costream;
106    detail::pythonbuf buffer;
107
108public:
109    scoped_ostream_redirect(
110            std::ostream &costream = std::cout,
111            object pyostream = module::import("sys").attr("stdout"))
112        : costream(costream), buffer(pyostream) {
113        old = costream.rdbuf(&buffer);
114    }
115
116    ~scoped_ostream_redirect() {
117        costream.rdbuf(old);
118    }
119
120    scoped_ostream_redirect(const scoped_ostream_redirect &) = delete;
121    scoped_ostream_redirect(scoped_ostream_redirect &&other) = default;
122    scoped_ostream_redirect &operator=(const scoped_ostream_redirect &) = delete;
123    scoped_ostream_redirect &operator=(scoped_ostream_redirect &&) = delete;
124};
125
126
127/** \rst
128    Like `scoped_ostream_redirect`, but redirects cerr by default. This class
129    is provided primary to make ``py::call_guard`` easier to make.
130
131    .. code-block:: cpp
132
133     m.def("noisy_func", &noisy_func,
134           py::call_guard<scoped_ostream_redirect,
135                          scoped_estream_redirect>());
136
137\endrst */
138class scoped_estream_redirect : public scoped_ostream_redirect {
139public:
140    scoped_estream_redirect(
141            std::ostream &costream = std::cerr,
142            object pyostream = module::import("sys").attr("stderr"))
143        : scoped_ostream_redirect(costream,pyostream) {}
144};
145
146
147NAMESPACE_BEGIN(detail)
148
149// Class to redirect output as a context manager. C++ backend.
150class OstreamRedirect {
151    bool do_stdout_;
152    bool do_stderr_;
153    std::unique_ptr<scoped_ostream_redirect> redirect_stdout;
154    std::unique_ptr<scoped_estream_redirect> redirect_stderr;
155
156public:
157    OstreamRedirect(bool do_stdout = true, bool do_stderr = true)
158        : do_stdout_(do_stdout), do_stderr_(do_stderr) {}
159
160    void enter() {
161        if (do_stdout_)
162            redirect_stdout.reset(new scoped_ostream_redirect());
163        if (do_stderr_)
164            redirect_stderr.reset(new scoped_estream_redirect());
165    }
166
167    void exit() {
168        redirect_stdout.reset();
169        redirect_stderr.reset();
170    }
171};
172
173NAMESPACE_END(detail)
174
175/** \rst
176    This is a helper function to add a C++ redirect context manager to Python
177    instead of using a C++ guard. To use it, add the following to your binding code:
178
179    .. code-block:: cpp
180
181        #include <pybind11/iostream.h>
182
183        ...
184
185        py::add_ostream_redirect(m, "ostream_redirect");
186
187    You now have a Python context manager that redirects your output:
188
189    .. code-block:: python
190
191        with m.ostream_redirect():
192            m.print_to_cout_function()
193
194    This manager can optionally be told which streams to operate on:
195
196    .. code-block:: python
197
198        with m.ostream_redirect(stdout=true, stderr=true):
199            m.noisy_function_with_error_printing()
200
201 \endrst */
202inline class_<detail::OstreamRedirect> add_ostream_redirect(module m, std::string name = "ostream_redirect") {
203    return class_<detail::OstreamRedirect>(m, name.c_str(), module_local())
204        .def(init<bool,bool>(), arg("stdout")=true, arg("stderr")=true)
205        .def("__enter__", &detail::OstreamRedirect::enter)
206        .def("__exit__", [](detail::OstreamRedirect &self_, args) { self_.exit(); });
207}
208
209NAMESPACE_END(PYBIND11_NAMESPACE)
210