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