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
20 PYBIND11_NAMESPACE_BEGIN(PYBIND11_NAMESPACE)
PYBIND11_NAMESPACE_BEGIN(detail)21 PYBIND11_NAMESPACE_BEGIN(detail)
22
23 // Buffer that writes to Python instead of C++
24 class pythonbuf : public std::streambuf {
25 private:
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) override {
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 // This function must be non-virtual to be called in a destructor. If the
42 // rare MSVC test failure shows up with this version, then this should be
43 // simplified to a fully qualified call.
44 int _sync() {
45 if (pbase() != pptr()) {
46 // This subtraction cannot be negative, so dropping the sign
47 str line(pbase(), static_cast<size_t>(pptr() - pbase()));
48
49 {
50 gil_scoped_acquire tmp;
51 pywrite(line);
52 pyflush();
53 }
54
55 setp(pbase(), epptr());
56 }
57 return 0;
58 }
59
60 int sync() override {
61 return _sync();
62 }
63
64 public:
65
66 pythonbuf(object pyostream, size_t buffer_size = 1024)
67 : buf_size(buffer_size),
68 d_buffer(new char[buf_size]),
69 pywrite(pyostream.attr("write")),
70 pyflush(pyostream.attr("flush")) {
71 setp(d_buffer.get(), d_buffer.get() + buf_size - 1);
72 }
73
74 pythonbuf(pythonbuf&&) = default;
75
76 /// Sync before destroy
77 ~pythonbuf() override {
78 _sync();
79 }
80 };
81
PYBIND11_NAMESPACE_END(detail)82 PYBIND11_NAMESPACE_END(detail)
83
84
85 /** \rst
86 This a move-only guard that redirects output.
87
88 .. code-block:: cpp
89
90 #include <pybind11/iostream.h>
91
92 ...
93
94 {
95 py::scoped_ostream_redirect output;
96 std::cout << "Hello, World!"; // Python stdout
97 } // <-- return std::cout to normal
98
99 You can explicitly pass the c++ stream and the python object,
100 for example to guard stderr instead.
101
102 .. code-block:: cpp
103
104 {
105 py::scoped_ostream_redirect output{std::cerr, py::module_::import("sys").attr("stderr")};
106 std::cerr << "Hello, World!";
107 }
108 \endrst */
109 class scoped_ostream_redirect {
110 protected:
111 std::streambuf *old;
112 std::ostream &costream;
113 detail::pythonbuf buffer;
114
115 public:
116 scoped_ostream_redirect(
117 std::ostream &costream = std::cout,
118 object pyostream = module_::import("sys").attr("stdout"))
119 : costream(costream), buffer(pyostream) {
120 old = costream.rdbuf(&buffer);
121 }
122
123 ~scoped_ostream_redirect() {
124 costream.rdbuf(old);
125 }
126
127 scoped_ostream_redirect(const scoped_ostream_redirect &) = delete;
128 scoped_ostream_redirect(scoped_ostream_redirect &&other) = default;
129 scoped_ostream_redirect &operator=(const scoped_ostream_redirect &) = delete;
130 scoped_ostream_redirect &operator=(scoped_ostream_redirect &&) = delete;
131 };
132
133
134 /** \rst
135 Like `scoped_ostream_redirect`, but redirects cerr by default. This class
136 is provided primary to make ``py::call_guard`` easier to make.
137
138 .. code-block:: cpp
139
140 m.def("noisy_func", &noisy_func,
141 py::call_guard<scoped_ostream_redirect,
142 scoped_estream_redirect>());
143
144 \endrst */
145 class scoped_estream_redirect : public scoped_ostream_redirect {
146 public:
147 scoped_estream_redirect(
148 std::ostream &costream = std::cerr,
149 object pyostream = module_::import("sys").attr("stderr"))
scoped_ostream_redirect(costream,pyostream)150 : scoped_ostream_redirect(costream,pyostream) {}
151 };
152
153
PYBIND11_NAMESPACE_BEGIN(detail)154 PYBIND11_NAMESPACE_BEGIN(detail)
155
156 // Class to redirect output as a context manager. C++ backend.
157 class OstreamRedirect {
158 bool do_stdout_;
159 bool do_stderr_;
160 std::unique_ptr<scoped_ostream_redirect> redirect_stdout;
161 std::unique_ptr<scoped_estream_redirect> redirect_stderr;
162
163 public:
164 OstreamRedirect(bool do_stdout = true, bool do_stderr = true)
165 : do_stdout_(do_stdout), do_stderr_(do_stderr) {}
166
167 void enter() {
168 if (do_stdout_)
169 redirect_stdout.reset(new scoped_ostream_redirect());
170 if (do_stderr_)
171 redirect_stderr.reset(new scoped_estream_redirect());
172 }
173
174 void exit() {
175 redirect_stdout.reset();
176 redirect_stderr.reset();
177 }
178 };
179
PYBIND11_NAMESPACE_END(detail)180 PYBIND11_NAMESPACE_END(detail)
181
182 /** \rst
183 This is a helper function to add a C++ redirect context manager to Python
184 instead of using a C++ guard. To use it, add the following to your binding code:
185
186 .. code-block:: cpp
187
188 #include <pybind11/iostream.h>
189
190 ...
191
192 py::add_ostream_redirect(m, "ostream_redirect");
193
194 You now have a Python context manager that redirects your output:
195
196 .. code-block:: python
197
198 with m.ostream_redirect():
199 m.print_to_cout_function()
200
201 This manager can optionally be told which streams to operate on:
202
203 .. code-block:: python
204
205 with m.ostream_redirect(stdout=true, stderr=true):
206 m.noisy_function_with_error_printing()
207
208 \endrst */
209 inline class_<detail::OstreamRedirect> add_ostream_redirect(module_ m, std::string name = "ostream_redirect") {
210 return class_<detail::OstreamRedirect>(m, name.c_str(), module_local())
211 .def(init<bool,bool>(), arg("stdout")=true, arg("stderr")=true)
212 .def("__enter__", &detail::OstreamRedirect::enter)
213 .def("__exit__", [](detail::OstreamRedirect &self_, args) { self_.exit(); });
214 }
215
216 PYBIND11_NAMESPACE_END(PYBIND11_NAMESPACE)
217