iostream.h revision 12391:ceeca8b41e4b 
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    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() ? 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
53public:
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
66NAMESPACE_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 */
93class scoped_ostream_redirect {
94protected:
95    std::streambuf *old;
96    std::ostream &costream;
97    detail::pythonbuf buffer;
98
99public:
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 */
129class scoped_estream_redirect : public scoped_ostream_redirect {
130public:
131    scoped_estream_redirect(
132            std::ostream &costream = std::cerr,
133            object pyostream = module::import("sys").attr("stderr"))
134        : scoped_ostream_redirect(costream,pyostream) {}
135};
136
137
138NAMESPACE_BEGIN(detail)
139
140// Class to redirect output as a context manager. C++ backend.
141class 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
147public:
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
164NAMESPACE_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 */
193inline 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
200NAMESPACE_END(PYBIND11_NAMESPACE)
201