HepMC3 event record library
iostream.h
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)
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  {
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 
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 */
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 */
146 public:
148  std::ostream &costream = std::cerr,
149  object pyostream = module_::import("sys").attr("stderr"))
150  : scoped_ostream_redirect(costream,pyostream) {}
151 };
152 
153 
154 PYBIND11_NAMESPACE_BEGIN(detail)
155 
156 // Class to redirect output as a context manager. C++ backend.
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 
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") {
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)
~pythonbuf() override
Sync before destroy.
Definition: iostream.h:77
Definition: pytypes.h:1341
Wrapper for Python extension modules.
Definition: pybind11.h:891
static module_ import(const char *name)
Import and return a module or throws error_already_set.
Definition: pybind11.h:941
Annotation that marks a class as local to the module:
Definition: attr.h:72
Definition: pytypes.h:935
Annotation for function names.
Definition: attr.h:36