include/pybind11/iostream.h

0001 /*
0002     pybind11/iostream.h -- Tools to assist with redirecting cout and cerr to Python
0003
0004     Copyright (c) 2017 Henry F. Schreiner
0005
0006     All rights reserved. Use of this source code is governed by a
0007     BSD-style license that can be found in the LICENSE file.
0008
0009     WARNING: The implementation in this file is NOT thread safe. Multiple
0010     threads writing to a redirected ostream concurrently cause data races
0011     and potentially buffer overflows. Therefore it is currently a requirement
0012     that all (possibly) concurrent redirected ostream writes are protected by
0013     a mutex.
0014     #HelpAppreciated: Work on iostream.h thread safety.
0015     For more background see the discussions under
0016     https://github.com/pybind/pybind11/pull/2982 and
0017     https://github.com/pybind/pybind11/pull/2995.
0018 */
0019
0020 #pragma once
0021
0022 #include "pybind11.h"
0023
0024 #include <algorithm>
0025 #include <cstring>
0026 #include <iostream>
0027 #include <iterator>
0028 #include <memory>
0029 #include <ostream>
0030 #include <streambuf>
0031 #include <string>
0032 #include <utility>
0033
0034 PYBIND11_NAMESPACE_BEGIN(PYBIND11_NAMESPACE)
0035 PYBIND11_NAMESPACE_BEGIN(detail)
0036
0037 // Buffer that writes to Python instead of C++
0038 class pythonbuf : public std::streambuf {
0039 private:
0040     using traits_type = std::streambuf::traits_type;
0041
0042     const size_t buf_size;
0043     std::unique_ptr<char[]> d_buffer;
0044     object pywrite;
0045     object pyflush;
0046
0047     int overflow(int c) override {
0048         if (!traits_type::eq_int_type(c, traits_type::eof())) {
0049             *pptr() = traits_type::to_char_type(c);
0050             pbump(1);
0051         }
0052         return sync() == 0 ? traits_type::not_eof(c) : traits_type::eof();
0053     }
0054
0055     // Computes how many bytes at the end of the buffer are part of an
0056     // incomplete sequence of UTF-8 bytes.
0057     // Precondition: pbase() < pptr()
0058     size_t utf8_remainder() const {
0059         const auto rbase = std::reverse_iterator<char *>(pbase());
0060         const auto rpptr = std::reverse_iterator<char *>(pptr());
0061         auto is_ascii = [](char c) { return (static_cast<unsigned char>(c) & 0x80) == 0x00; };
0062         auto is_leading = [](char c) { return (static_cast<unsigned char>(c) & 0xC0) == 0xC0; };
0063         auto is_leading_2b = [](char c) { return static_cast<unsigned char>(c) <= 0xDF; };
0064         auto is_leading_3b = [](char c) { return static_cast<unsigned char>(c) <= 0xEF; };
0065         // If the last character is ASCII, there are no incomplete code points
0066         if (is_ascii(*rpptr)) {
0067             return 0;
0068         }
0069         // Otherwise, work back from the end of the buffer and find the first
0070         // UTF-8 leading byte
0071         const auto rpend = rbase - rpptr >= 3 ? rpptr + 3 : rbase;
0072         const auto leading = std::find_if(rpptr, rpend, is_leading);
0073         if (leading == rbase) {
0074             return 0;
0075         }
0076         const auto dist = static_cast<size_t>(leading - rpptr);
0077         size_t remainder = 0;
0078
0079         if (dist == 0) {
0080             remainder = 1; // 1-byte code point is impossible
0081         } else if (dist == 1) {
0082             remainder = is_leading_2b(*leading) ? 0 : dist + 1;
0083         } else if (dist == 2) {
0084             remainder = is_leading_3b(*leading) ? 0 : dist + 1;
0085         }
0086         // else if (dist >= 3), at least 4 bytes before encountering an UTF-8
0087         // leading byte, either no remainder or invalid UTF-8.
0088         // Invalid UTF-8 will cause an exception later when converting
0089         // to a Python string, so that's not handled here.
0090         return remainder;
0091     }
0092
0093     // This function must be non-virtual to be called in a destructor.
0094     int _sync() {
0095         if (pbase() != pptr()) { // If buffer is not empty
0096             gil_scoped_acquire tmp;
0097             // This subtraction cannot be negative, so dropping the sign.
0098             auto size = static_cast<size_t>(pptr() - pbase());
0099             size_t remainder = utf8_remainder();
0100
0101             if (size > remainder) {
0102                 str line(pbase(), size - remainder);
0103                 pywrite(std::move(line));
0104                 pyflush();
0105             }
0106
0107             // Copy the remainder at the end of the buffer to the beginning:
0108             if (remainder > 0) {
0109                 std::memmove(pbase(), pptr() - remainder, remainder);
0110             }
0111             setp(pbase(), epptr());
0112             pbump(static_cast<int>(remainder));
0113         }
0114         return 0;
0115     }
0116
0117     int sync() override { return _sync(); }
0118
0119 public:
0120     explicit pythonbuf(const object &pyostream, size_t buffer_size = 1024)
0121         : buf_size(buffer_size), d_buffer(new char[buf_size]), pywrite(pyostream.attr("write")),
0122           pyflush(pyostream.attr("flush")) {
0123         setp(d_buffer.get(), d_buffer.get() + buf_size - 1);
0124     }
0125
0126     pythonbuf(pythonbuf &&) = default;
0127
0128     /// Sync before destroy
0129     ~pythonbuf() override { _sync(); }
0130 };
0131
0132 PYBIND11_NAMESPACE_END(detail)
0133
0134 /** \rst
0135     This a move-only guard that redirects output.
0136
0137     .. code-block:: cpp
0138
0139         #include <pybind11/iostream.h>
0140
0141         ...
0142
0143         {
0144             py::scoped_ostream_redirect output;
0145             std::cout << "Hello, World!"; // Python stdout
0146         } // <-- return std::cout to normal
0147
0148     You can explicitly pass the c++ stream and the python object,
0149     for example to guard stderr instead.
0150
0151     .. code-block:: cpp
0152
0153         {
0154             py::scoped_ostream_redirect output{
0155                 std::cerr, py::module::import("sys").attr("stderr")};
0156             std::cout << "Hello, World!";
0157         }
0158  \endrst */
0159 class scoped_ostream_redirect {
0160 protected:
0161     std::streambuf *old;
0162     std::ostream &costream;
0163     detail::pythonbuf buffer;
0164
0165 public:
0166     explicit scoped_ostream_redirect(std::ostream &costream = std::cout,
0167                                      const object &pyostream
0168                                      = module_::import("sys").attr("stdout"))
0169         : costream(costream), buffer(pyostream) {
0170         old = costream.rdbuf(&buffer);
0171     }
0172
0173     ~scoped_ostream_redirect() { costream.rdbuf(old); }
0174
0175     scoped_ostream_redirect(const scoped_ostream_redirect &) = delete;
0176     scoped_ostream_redirect(scoped_ostream_redirect &&other) = default;
0177     scoped_ostream_redirect &operator=(const scoped_ostream_redirect &) = delete;
0178     scoped_ostream_redirect &operator=(scoped_ostream_redirect &&) = delete;
0179 };
0180
0181 /** \rst
0182     Like `scoped_ostream_redirect`, but redirects cerr by default. This class
0183     is provided primary to make ``py::call_guard`` easier to make.
0184
0185     .. code-block:: cpp
0186
0187      m.def("noisy_func", &noisy_func,
0188            py::call_guard<scoped_ostream_redirect,
0189                           scoped_estream_redirect>());
0190
0191 \endrst */
0192 class scoped_estream_redirect : public scoped_ostream_redirect {
0193 public:
0194     explicit scoped_estream_redirect(std::ostream &costream = std::cerr,
0195                                      const object &pyostream
0196                                      = module_::import("sys").attr("stderr"))
0197         : scoped_ostream_redirect(costream, pyostream) {}
0198 };
0199
0200 PYBIND11_NAMESPACE_BEGIN(detail)
0201
0202 // Class to redirect output as a context manager. C++ backend.
0203 class OstreamRedirect {
0204     bool do_stdout_;
0205     bool do_stderr_;
0206     std::unique_ptr<scoped_ostream_redirect> redirect_stdout;
0207     std::unique_ptr<scoped_estream_redirect> redirect_stderr;
0208
0209 public:
0210     explicit OstreamRedirect(bool do_stdout = true, bool do_stderr = true)
0211         : do_stdout_(do_stdout), do_stderr_(do_stderr) {}
0212
0213     void enter() {
0214         if (do_stdout_) {
0215             redirect_stdout.reset(new scoped_ostream_redirect());
0216         }
0217         if (do_stderr_) {
0218             redirect_stderr.reset(new scoped_estream_redirect());
0219         }
0220     }
0221
0222     void exit() {
0223         redirect_stdout.reset();
0224         redirect_stderr.reset();
0225     }
0226 };
0227
0228 PYBIND11_NAMESPACE_END(detail)
0229
0230 /** \rst
0231     This is a helper function to add a C++ redirect context manager to Python
0232     instead of using a C++ guard. To use it, add the following to your binding code:
0233
0234     .. code-block:: cpp
0235
0236         #include <pybind11/iostream.h>
0237
0238         ...
0239
0240         py::add_ostream_redirect(m, "ostream_redirect");
0241
0242     You now have a Python context manager that redirects your output:
0243
0244     .. code-block:: python
0245
0246         with m.ostream_redirect():
0247             m.print_to_cout_function()
0248
0249     This manager can optionally be told which streams to operate on:
0250
0251     .. code-block:: python
0252
0253         with m.ostream_redirect(stdout=true, stderr=true):
0254             m.noisy_function_with_error_printing()
0255
0256  \endrst */
0257 inline class_<detail::OstreamRedirect>
0258 add_ostream_redirect(module_ m, const std::string &name = "ostream_redirect") {
0259     return class_<detail::OstreamRedirect>(std::move(m), name.c_str(), module_local())
0260         .def(init<bool, bool>(), arg("stdout") = true, arg("stderr") = true)
0261         .def("__enter__", &detail::OstreamRedirect::enter)
0262         .def("__exit__", [](detail::OstreamRedirect &self_, const args &) { self_.exit(); });
0263 }
0264
0265 PYBIND11_NAMESPACE_END(PYBIND11_NAMESPACE)