1 /*
2 Copyright (c) 2016 Wenzel Jakob <wenzel.jakob@epfl.ch>
3
4 All rights reserved. Use of this source code is governed by a
5 BSD-style license that can be found in the LICENSE file.
6 */
7
8
9
10
11 #pragma once
12
13 #include "pybind11.h"
14
15 #include <streambuf>
16 #include <ostream>
17 #include <string>
18 #include <memory>
19 #include <iostream>
20
21 NAMESPACE_BEGIN(PYBIND11_NAMESPACE)
NAMESPACE_BEGIN(detail)22 NAMESPACE_BEGIN(detail)
23
24
25 class pythonbuf : public std::streambuf {
26 private:
27 using traits_type = std::streambuf::traits_type;
28
29 char d_buffer[1024];
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
44 str line(pbase(), static_cast<size_t>(pptr() - pbase()));
45
46 pywrite(line);
47 pyflush();
48
49 setp(pbase(), epptr());
50 }
51 return 0;
52 }
53
54 public:
55 pythonbuf(object pyostream)
56 : pywrite(pyostream.attr("write")),
57 pyflush(pyostream.attr("flush")) {
58 setp(d_buffer, d_buffer + sizeof(d_buffer) - 1);
59 }
60
61
62 ~pythonbuf() {
63 sync();
64 }
65 };
66
NAMESPACE_END(detail)67 NAMESPACE_END(detail)
68
69
70
71 class scoped_ostream_redirect {
72 protected:
73 std::streambuf *old;
74 std::ostream &costream;
75 detail::pythonbuf buffer;
76
77 public:
78 scoped_ostream_redirect(
79 std::ostream &costream = std::cout,
80 object pyostream = module::import("sys").attr("stdout"))
81 : costream(costream), buffer(pyostream) {
82 old = costream.rdbuf(&buffer);
83 }
84
85 ~scoped_ostream_redirect() {
86 costream.rdbuf(old);
87 }
88
89 scoped_ostream_redirect(const scoped_ostream_redirect &) = delete;
90 scoped_ostream_redirect(scoped_ostream_redirect &&other) = default;
91 scoped_ostream_redirect &operator=(const scoped_ostream_redirect &) = delete;
92 scoped_ostream_redirect &operator=(scoped_ostream_redirect &&) = delete;
93 };
94
95
96
97 class scoped_estream_redirect : public scoped_ostream_redirect {
98 public:
99 scoped_estream_redirect(
100 std::ostream &costream = std::cerr,
101 object pyostream = module::import("sys").attr("stderr"))
scoped_ostream_redirect(costream,pyostream)102 : scoped_ostream_redirect(costream,pyostream) {}
103 };
104
105
NAMESPACE_BEGIN(detail)106 NAMESPACE_BEGIN(detail)
107
108
109 class OstreamRedirect {
110 bool do_stdout_;
111 bool do_stderr_;
112 std::unique_ptr<scoped_ostream_redirect> redirect_stdout;
113 std::unique_ptr<scoped_estream_redirect> redirect_stderr;
114
115 public:
116 OstreamRedirect(bool do_stdout = true, bool do_stderr = true)
117 : do_stdout_(do_stdout), do_stderr_(do_stderr) {}
118
119 void enter() {
120 if (do_stdout_)
121 redirect_stdout.reset(new scoped_ostream_redirect());
122 if (do_stderr_)
123 redirect_stderr.reset(new scoped_estream_redirect());
124 }
125
126 void exit() {
127 redirect_stdout.reset();
128 redirect_stderr.reset();
129 }
130 };
131
NAMESPACE_END(detail)132 NAMESPACE_END(detail)
133
134
135 inline class_<detail::OstreamRedirect> add_ostream_redirect(module m, std::string name = "ostream_redirect") {
136 return class_<detail::OstreamRedirect>(m, name.c_str(), module_local())
137 .def(init<bool,bool>(), arg("stdout")=true, arg("stderr")=true)
138 .def("__enter__", &detail::OstreamRedirect::enter)
139 .def("__exit__", [](detail::OstreamRedirect &self_, args) { self_.exit(); });
140 }
141
142 NAMESPACE_END(PYBIND11_NAMESPACE)
143