ext/pybind11/include/pybind11/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 NAMESPACE_BEGIN(PYBIND11_NAMESPACE)
  21 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) {
  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             // This subtraction cannot be negative, so dropping the sign
  44             str line(pbase(), static_cast<size_t>(pptr() - pbase()));
  45
  46             {
  47                 gil_scoped_acquire tmp;
  48                 pywrite(line);
  49                 pyflush();
  50             }
  51
  52             setp(pbase(), epptr());
  53         }
  54         return 0;
  55     }
  56
  57 public:
  58
  59     pythonbuf(object pyostream, size_t buffer_size = 1024)
  60         : buf_size(buffer_size),
  61           d_buffer(new char[buf_size]),
  62           pywrite(pyostream.attr("write")),
  63           pyflush(pyostream.attr("flush")) {
  64         setp(d_buffer.get(), d_buffer.get() + buf_size - 1);
  65     }
  66
  67     pythonbuf(pythonbuf&&) = default;
  68
  69     /// Sync before destroy
  70     ~pythonbuf() {
  71         sync();
  72     }
  73 };
  74
  75 NAMESPACE_END(detail)
  76
  77
  78 /** \rst
  79     This a move-only guard that redirects output.
  80
  81     .. code-block:: cpp
  82
  83         #include <pybind11/iostream.h>
  84
  85         ...
  86
  87         {
  88             py::scoped_ostream_redirect output;
  89             std::cout << "Hello, World!"; // Python stdout
  90         } // <-- return std::cout to normal
  91
  92     You can explicitly pass the c++ stream and the python object,
  93     for example to guard stderr instead.
  94
  95     .. code-block:: cpp
  96
  97         {
  98             py::scoped_ostream_redirect output{std::cerr, py::module::import("sys").attr("stderr")};
  99             std::cerr << "Hello, World!";
 100         }
 101  \endrst */
 102 class scoped_ostream_redirect {
 103 protected:
 104     std::streambuf *old;
 105     std::ostream &costream;
 106     detail::pythonbuf buffer;
 107
 108 public:
 109     scoped_ostream_redirect(
 110             std::ostream &costream = std::cout,
 111             object pyostream = module::import("sys").attr("stdout"))
 112         : costream(costream), buffer(pyostream) {
 113         old = costream.rdbuf(&buffer);
 114     }
 115
 116     ~scoped_ostream_redirect() {
 117         costream.rdbuf(old);
 118     }
 119
 120     scoped_ostream_redirect(const scoped_ostream_redirect &) = delete;
 121     scoped_ostream_redirect(scoped_ostream_redirect &&other) = default;
 122     scoped_ostream_redirect &operator=(const scoped_ostream_redirect &) = delete;
 123     scoped_ostream_redirect &operator=(scoped_ostream_redirect &&) = delete;
 124 };
 125
 126
 127 /** \rst
 128     Like `scoped_ostream_redirect`, but redirects cerr by default. This class
 129     is provided primary to make ``py::call_guard`` easier to make.
 130
 131     .. code-block:: cpp
 132
 133      m.def("noisy_func", &noisy_func,
 134            py::call_guard<scoped_ostream_redirect,
 135                           scoped_estream_redirect>());
 136
 137 \endrst */
 138 class scoped_estream_redirect : public scoped_ostream_redirect {
 139 public:
 140     scoped_estream_redirect(
 141             std::ostream &costream = std::cerr,
 142             object pyostream = module::import("sys").attr("stderr"))
 143         : scoped_ostream_redirect(costream,pyostream) {}
 144 };
 145
 146
 147 NAMESPACE_BEGIN(detail)
 148
 149 // Class to redirect output as a context manager. C++ backend.
 150 class OstreamRedirect {
 151     bool do_stdout_;
 152     bool do_stderr_;
 153     std::unique_ptr<scoped_ostream_redirect> redirect_stdout;
 154     std::unique_ptr<scoped_estream_redirect> redirect_stderr;
 155
 156 public:
 157     OstreamRedirect(bool do_stdout = true, bool do_stderr = true)
 158         : do_stdout_(do_stdout), do_stderr_(do_stderr) {}
 159
 160     void enter() {
 161         if (do_stdout_)
 162             redirect_stdout.reset(new scoped_ostream_redirect());
 163         if (do_stderr_)
 164             redirect_stderr.reset(new scoped_estream_redirect());
 165     }
 166
 167     void exit() {
 168         redirect_stdout.reset();
 169         redirect_stderr.reset();
 170     }
 171 };
 172
 173 NAMESPACE_END(detail)
 174
 175 /** \rst
 176     This is a helper function to add a C++ redirect context manager to Python
 177     instead of using a C++ guard. To use it, add the following to your binding code:
 178
 179     .. code-block:: cpp
 180
 181         #include <pybind11/iostream.h>
 182
 183         ...
 184
 185         py::add_ostream_redirect(m, "ostream_redirect");
 186
 187     You now have a Python context manager that redirects your output:
 188
 189     .. code-block:: python
 190
 191         with m.ostream_redirect():
 192             m.print_to_cout_function()
 193
 194     This manager can optionally be told which streams to operate on:
 195
 196     .. code-block:: python
 197
 198         with m.ostream_redirect(stdout=true, stderr=true):
 199             m.noisy_function_with_error_printing()
 200
 201  \endrst */
 202 inline class_<detail::OstreamRedirect> add_ostream_redirect(module m, std::string name = "ostream_redirect") {
 203     return class_<detail::OstreamRedirect>(m, name.c_str(), module_local())
 204         .def(init<bool,bool>(), arg("stdout")=true, arg("stderr")=true)
 205         .def("__enter__", &detail::OstreamRedirect::enter)
 206         .def("__exit__", [](detail::OstreamRedirect &self_, args) { self_.exit(); });
 207 }
 208
 209 NAMESPACE_END(PYBIND11_NAMESPACE)