iostream.h 5.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220
  1. /*
  2. pybind11/iostream.h -- Tools to assist with redirecting cout and cerr to Python
  3. Copyright (c) 2017 Henry F. Schreiner
  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. #pragma once
  8. #include "pybind11.h"
  9. #include <streambuf>
  10. #include <ostream>
  11. #include <string>
  12. #include <memory>
  13. #include <iostream>
  14. PYBIND11_NAMESPACE_BEGIN(PYBIND11_NAMESPACE)
  15. PYBIND11_NAMESPACE_BEGIN(detail)
  16. // Buffer that writes to Python instead of C++
  17. class pythonbuf : public std::streambuf {
  18. private:
  19. using traits_type = std::streambuf::traits_type;
  20. const size_t buf_size;
  21. std::unique_ptr<char[]> d_buffer;
  22. object pywrite;
  23. object pyflush;
  24. int overflow(int c) override {
  25. if (!traits_type::eq_int_type(c, traits_type::eof())) {
  26. *pptr() = traits_type::to_char_type(c);
  27. pbump(1);
  28. }
  29. return sync() == 0 ? traits_type::not_eof(c) : traits_type::eof();
  30. }
  31. // This function must be non-virtual to be called in a destructor. If the
  32. // rare MSVC test failure shows up with this version, then this should be
  33. // simplified to a fully qualified call.
  34. int _sync() {
  35. if (pbase() != pptr()) {
  36. {
  37. gil_scoped_acquire tmp;
  38. // This subtraction cannot be negative, so dropping the sign.
  39. str line(pbase(), static_cast<size_t>(pptr() - pbase()));
  40. pywrite(line);
  41. pyflush();
  42. // Placed inside gil_scoped_aquire as a mutex to avoid a race
  43. setp(pbase(), epptr());
  44. }
  45. }
  46. return 0;
  47. }
  48. int sync() override {
  49. return _sync();
  50. }
  51. public:
  52. pythonbuf(object pyostream, size_t buffer_size = 1024)
  53. : buf_size(buffer_size),
  54. d_buffer(new char[buf_size]),
  55. pywrite(pyostream.attr("write")),
  56. pyflush(pyostream.attr("flush")) {
  57. setp(d_buffer.get(), d_buffer.get() + buf_size - 1);
  58. }
  59. pythonbuf(pythonbuf&&) = default;
  60. /// Sync before destroy
  61. ~pythonbuf() override {
  62. _sync();
  63. }
  64. };
  65. PYBIND11_NAMESPACE_END(detail)
  66. /** \rst
  67. This a move-only guard that redirects output.
  68. .. code-block:: cpp
  69. #include <pybind11/iostream.h>
  70. ...
  71. {
  72. py::scoped_ostream_redirect output;
  73. std::cout << "Hello, World!"; // Python stdout
  74. } // <-- return std::cout to normal
  75. You can explicitly pass the c++ stream and the python object,
  76. for example to guard stderr instead.
  77. .. code-block:: cpp
  78. {
  79. py::scoped_ostream_redirect output{std::cerr, py::module::import("sys").attr("stderr")};
  80. std::cout << "Hello, World!";
  81. }
  82. \endrst */
  83. class scoped_ostream_redirect {
  84. protected:
  85. std::streambuf *old;
  86. std::ostream &costream;
  87. detail::pythonbuf buffer;
  88. public:
  89. scoped_ostream_redirect(
  90. std::ostream &costream = std::cout,
  91. object pyostream = module_::import("sys").attr("stdout"))
  92. : costream(costream), buffer(pyostream) {
  93. old = costream.rdbuf(&buffer);
  94. }
  95. ~scoped_ostream_redirect() {
  96. costream.rdbuf(old);
  97. }
  98. scoped_ostream_redirect(const scoped_ostream_redirect &) = delete;
  99. scoped_ostream_redirect(scoped_ostream_redirect &&other) = default;
  100. scoped_ostream_redirect &operator=(const scoped_ostream_redirect &) = delete;
  101. scoped_ostream_redirect &operator=(scoped_ostream_redirect &&) = delete;
  102. };
  103. /** \rst
  104. Like `scoped_ostream_redirect`, but redirects cerr by default. This class
  105. is provided primary to make ``py::call_guard`` easier to make.
  106. .. code-block:: cpp
  107. m.def("noisy_func", &noisy_func,
  108. py::call_guard<scoped_ostream_redirect,
  109. scoped_estream_redirect>());
  110. \endrst */
  111. class scoped_estream_redirect : public scoped_ostream_redirect {
  112. public:
  113. scoped_estream_redirect(
  114. std::ostream &costream = std::cerr,
  115. object pyostream = module_::import("sys").attr("stderr"))
  116. : scoped_ostream_redirect(costream,pyostream) {}
  117. };
  118. PYBIND11_NAMESPACE_BEGIN(detail)
  119. // Class to redirect output as a context manager. C++ backend.
  120. class OstreamRedirect {
  121. bool do_stdout_;
  122. bool do_stderr_;
  123. std::unique_ptr<scoped_ostream_redirect> redirect_stdout;
  124. std::unique_ptr<scoped_estream_redirect> redirect_stderr;
  125. public:
  126. OstreamRedirect(bool do_stdout = true, bool do_stderr = true)
  127. : do_stdout_(do_stdout), do_stderr_(do_stderr) {}
  128. void enter() {
  129. if (do_stdout_)
  130. redirect_stdout.reset(new scoped_ostream_redirect());
  131. if (do_stderr_)
  132. redirect_stderr.reset(new scoped_estream_redirect());
  133. }
  134. void exit() {
  135. redirect_stdout.reset();
  136. redirect_stderr.reset();
  137. }
  138. };
  139. PYBIND11_NAMESPACE_END(detail)
  140. /** \rst
  141. This is a helper function to add a C++ redirect context manager to Python
  142. instead of using a C++ guard. To use it, add the following to your binding code:
  143. .. code-block:: cpp
  144. #include <pybind11/iostream.h>
  145. ...
  146. py::add_ostream_redirect(m, "ostream_redirect");
  147. You now have a Python context manager that redirects your output:
  148. .. code-block:: python
  149. with m.ostream_redirect():
  150. m.print_to_cout_function()
  151. This manager can optionally be told which streams to operate on:
  152. .. code-block:: python
  153. with m.ostream_redirect(stdout=true, stderr=true):
  154. m.noisy_function_with_error_printing()
  155. \endrst */
  156. inline class_<detail::OstreamRedirect> add_ostream_redirect(module_ m, std::string name = "ostream_redirect") {
  157. return class_<detail::OstreamRedirect>(m, name.c_str(), module_local())
  158. .def(init<bool,bool>(), arg("stdout")=true, arg("stderr")=true)
  159. .def("__enter__", &detail::OstreamRedirect::enter)
  160. .def("__exit__", [](detail::OstreamRedirect &self_, args) { self_.exit(); });
  161. }
  162. PYBIND11_NAMESPACE_END(PYBIND11_NAMESPACE)