iostream.h
Go to the documentation of this file.
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  WARNING: The implementation in this file is NOT thread safe. Multiple
10  threads writing to a redirected ostream concurrently cause data races
11  and potentially buffer overflows. Therefore it is currently a requirement
12  that all (possibly) concurrent redirected ostream writes are protected by
13  a mutex.
14  #HelpAppreciated: Work on iostream.h thread safety.
15  For more background see the discussions under
16  https://github.com/pybind/pybind11/pull/2982 and
17  https://github.com/pybind/pybind11/pull/2995.
18 */
19 
20 #pragma once
21 
22 #include "pybind11.h"
23 
24 #include <algorithm>
25 #include <cstring>
26 #include <iostream>
27 #include <iterator>
28 #include <memory>
29 #include <ostream>
30 #include <streambuf>
31 #include <string>
32 #include <utility>
33 
36 
37 // Buffer that writes to Python instead of C++
38 class pythonbuf : public std::streambuf {
39 private:
40  using traits_type = std::streambuf::traits_type;
41 
42  const size_t buf_size;
43  std::unique_ptr<char[]> d_buffer;
44  object pywrite;
45  object pyflush;
46 
47  int overflow(int c) override {
48  if (!traits_type::eq_int_type(c, traits_type::eof())) {
49  *pptr() = traits_type::to_char_type(c);
50  pbump(1);
51  }
52  return sync() == 0 ? traits_type::not_eof(c) : traits_type::eof();
53  }
54 
55  // Computes how many bytes at the end of the buffer are part of an
56  // incomplete sequence of UTF-8 bytes.
57  // Precondition: pbase() < pptr()
58  size_t utf8_remainder() const {
59  const auto rbase = std::reverse_iterator<char *>(pbase());
60  const auto rpptr = std::reverse_iterator<char *>(pptr());
61  auto is_ascii = [](char c) { return (static_cast<unsigned char>(c) & 0x80) == 0x00; };
62  auto is_leading = [](char c) { return (static_cast<unsigned char>(c) & 0xC0) == 0xC0; };
63  auto is_leading_2b = [](char c) { return static_cast<unsigned char>(c) <= 0xDF; };
64  auto is_leading_3b = [](char c) { return static_cast<unsigned char>(c) <= 0xEF; };
65  // If the last character is ASCII, there are no incomplete code points
66  if (is_ascii(*rpptr)) {
67  return 0;
68  }
69  // Otherwise, work back from the end of the buffer and find the first
70  // UTF-8 leading byte
71  const auto rpend = rbase - rpptr >= 3 ? rpptr + 3 : rbase;
72  const auto leading = std::find_if(rpptr, rpend, is_leading);
73  if (leading == rbase) {
74  return 0;
75  }
76  const auto dist = static_cast<size_t>(leading - rpptr);
77  size_t remainder = 0;
78 
79  if (dist == 0) {
80  remainder = 1; // 1-byte code point is impossible
81  } else if (dist == 1) {
82  remainder = is_leading_2b(*leading) ? 0 : dist + 1;
83  } else if (dist == 2) {
84  remainder = is_leading_3b(*leading) ? 0 : dist + 1;
85  }
86  // else if (dist >= 3), at least 4 bytes before encountering an UTF-8
87  // leading byte, either no remainder or invalid UTF-8.
88  // Invalid UTF-8 will cause an exception later when converting
89  // to a Python string, so that's not handled here.
90  return remainder;
91  }
92 
93  // This function must be non-virtual to be called in a destructor.
94  int _sync() {
95  if (pbase() != pptr()) { // If buffer is not empty
97  // This subtraction cannot be negative, so dropping the sign.
98  auto size = static_cast<size_t>(pptr() - pbase());
99  size_t remainder = utf8_remainder();
100 
101  if (size > remainder) {
102  str line(pbase(), size - remainder);
103  pywrite(std::move(line));
104  pyflush();
105  }
106 
107  // Copy the remainder at the end of the buffer to the beginning:
108  if (remainder > 0) {
109  std::memmove(pbase(), pptr() - remainder, remainder);
110  }
111  setp(pbase(), epptr());
112  pbump(static_cast<int>(remainder));
113  }
114  return 0;
115  }
116 
117  int sync() override { return _sync(); }
118 
119 public:
120  explicit pythonbuf(const object &pyostream, size_t buffer_size = 1024)
121  : buf_size(buffer_size), d_buffer(new char[buf_size]), pywrite(pyostream.attr("write")),
122  pyflush(pyostream.attr("flush")) {
123  setp(d_buffer.get(), d_buffer.get() + buf_size - 1);
124  }
125 
126  pythonbuf(pythonbuf &&) = default;
127 
129  ~pythonbuf() override { _sync(); }
130 };
131 
133 
134 
160 protected:
161  std::streambuf *old;
162  std::ostream &costream;
163  detail::pythonbuf buffer;
164 
165 public:
166  explicit scoped_ostream_redirect(std::ostream &costream = std::cout,
167  const object &pyostream
168  = module_::import("sys").attr("stdout"))
169  : costream(costream), buffer(pyostream) {
170  old = costream.rdbuf(&buffer);
171  }
172 
173  ~scoped_ostream_redirect() { costream.rdbuf(old); }
174 
177  scoped_ostream_redirect &operator=(const scoped_ostream_redirect &) = delete;
178  scoped_ostream_redirect &operator=(scoped_ostream_redirect &&) = delete;
179 };
180 
193 public:
194  explicit scoped_estream_redirect(std::ostream &costream = std::cerr,
195  const object &pyostream
196  = module_::import("sys").attr("stderr"))
197  : scoped_ostream_redirect(costream, pyostream) {}
198 };
199 
201 
202 // Class to redirect output as a context manager. C++ backend.
206  std::unique_ptr<scoped_ostream_redirect> redirect_stdout;
207  std::unique_ptr<scoped_estream_redirect> redirect_stderr;
208 
209 public:
210  explicit OstreamRedirect(bool do_stdout = true, bool do_stderr = true)
211  : do_stdout_(do_stdout), do_stderr_(do_stderr) {}
212 
213  void enter() {
214  if (do_stdout_) {
215  redirect_stdout.reset(new scoped_ostream_redirect());
216  }
217  if (do_stderr_) {
218  redirect_stderr.reset(new scoped_estream_redirect());
219  }
220  }
221 
222  void exit() {
223  redirect_stdout.reset();
224  redirect_stderr.reset();
225  }
226 };
227 
229 
230 
258 add_ostream_redirect(module_ m, const std::string &name = "ostream_redirect") {
259  return class_<detail::OstreamRedirect>(std::move(m), name.c_str(), module_local())
260  .def(init<bool, bool>(), arg("stdout") = true, arg("stderr") = true)
261  .def("__enter__", &detail::OstreamRedirect::enter)
262  .def("__exit__", [](detail::OstreamRedirect &self_, const args &) { self_.exit(); });
263 }
264 
object pywrite
Definition: iostream.h:44
Matrix3f m
std::streambuf::traits_type traits_type
Definition: iostream.h:40
~pythonbuf() override
Sync before destroy.
Definition: iostream.h:129
int sync() override
Definition: iostream.h:117
int _sync()
Definition: iostream.h:94
scoped_estream_redirect(std::ostream &costream=std::cerr, const object &pyostream=module_::import("sys").attr("stderr"))
Definition: iostream.h:194
Definition: pytypes.h:2012
Scalar Scalar * c
Definition: benchVecAdd.cpp:17
class_ & def(const char *name_, Func &&f, const Extra &...extra)
Definition: pybind11.h:1557
Wrapper for Python extension modules.
Definition: pybind11.h:1136
class_< detail::OstreamRedirect > add_ostream_redirect(module_ m, const std::string &name="ostream_redirect")
Definition: iostream.h:258
int overflow(int c) override
Definition: iostream.h:47
Definition: BFloat16.h:88
pythonbuf(const object &pyostream, size_t buffer_size=1024)
Definition: iostream.h:120
const size_t buf_size
Definition: iostream.h:42
detail::pythonbuf buffer
Definition: iostream.h:163
std::streambuf * old
Definition: iostream.h:161
Scalar Scalar int size
Definition: benchVecAdd.cpp:17
Annotation that marks a class as local to the module:
Definition: attr.h:110
Definition: pytypes.h:1403
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const ArgReturnType arg() const
std::unique_ptr< scoped_estream_redirect > redirect_stderr
Definition: iostream.h:207
scoped_ostream_redirect(std::ostream &costream=std::cout, const object &pyostream=module_::import("sys").attr("stdout"))
Definition: iostream.h:166
std::unique_ptr< char[]> d_buffer
Definition: iostream.h:43
std::unique_ptr< scoped_ostream_redirect > redirect_stdout
Definition: iostream.h:206
static module_ import(const char *name)
Import and return a module or throws error_already_set.
Definition: pybind11.h:1194
OstreamRedirect(bool do_stdout=true, bool do_stderr=true)
Definition: iostream.h:210
Annotation for function names.
Definition: attr.h:48
object pyflush
Definition: iostream.h:45
void enter()
Definition: iostream.h:213
std::ostream & costream
Definition: iostream.h:162
size_t utf8_remainder() const
Definition: iostream.h:58
#define PYBIND11_NAMESPACE_END(name)
#define PYBIND11_NAMESPACE_BEGIN(name)


gtsam
Author(s):
autogenerated on Tue Jul 4 2023 02:34:25