iostream.h revision 12391
112391Sjason@lowepower.com/* 212391Sjason@lowepower.com pybind11/iostream.h -- Tools to assist with redirecting cout and cerr to Python 312391Sjason@lowepower.com 412391Sjason@lowepower.com Copyright (c) 2017 Henry F. Schreiner 512391Sjason@lowepower.com 612391Sjason@lowepower.com All rights reserved. Use of this source code is governed by a 712391Sjason@lowepower.com BSD-style license that can be found in the LICENSE file. 812391Sjason@lowepower.com*/ 912391Sjason@lowepower.com 1012391Sjason@lowepower.com#pragma once 1112391Sjason@lowepower.com 1212391Sjason@lowepower.com#include "pybind11.h" 1312391Sjason@lowepower.com 1412391Sjason@lowepower.com#include <streambuf> 1512391Sjason@lowepower.com#include <ostream> 1612391Sjason@lowepower.com#include <string> 1712391Sjason@lowepower.com#include <memory> 1812391Sjason@lowepower.com#include <iostream> 1912391Sjason@lowepower.com 2012391Sjason@lowepower.comNAMESPACE_BEGIN(PYBIND11_NAMESPACE) 2112391Sjason@lowepower.comNAMESPACE_BEGIN(detail) 2212391Sjason@lowepower.com 2312391Sjason@lowepower.com// Buffer that writes to Python instead of C++ 2412391Sjason@lowepower.comclass pythonbuf : public std::streambuf { 2512391Sjason@lowepower.comprivate: 2612391Sjason@lowepower.com using traits_type = std::streambuf::traits_type; 2712391Sjason@lowepower.com 2812391Sjason@lowepower.com char d_buffer[1024]; 2912391Sjason@lowepower.com object pywrite; 3012391Sjason@lowepower.com object pyflush; 3112391Sjason@lowepower.com 3212391Sjason@lowepower.com int overflow(int c) { 3312391Sjason@lowepower.com if (!traits_type::eq_int_type(c, traits_type::eof())) { 3412391Sjason@lowepower.com *pptr() = traits_type::to_char_type(c); 3512391Sjason@lowepower.com pbump(1); 3612391Sjason@lowepower.com } 3712391Sjason@lowepower.com return sync() ? traits_type::not_eof(c) : traits_type::eof(); 3812391Sjason@lowepower.com } 3912391Sjason@lowepower.com 4012391Sjason@lowepower.com int sync() { 4112391Sjason@lowepower.com if (pbase() != pptr()) { 4212391Sjason@lowepower.com // This subtraction cannot be negative, so dropping the sign 4312391Sjason@lowepower.com str line(pbase(), static_cast<size_t>(pptr() - pbase())); 4412391Sjason@lowepower.com 4512391Sjason@lowepower.com pywrite(line); 4612391Sjason@lowepower.com pyflush(); 4712391Sjason@lowepower.com 4812391Sjason@lowepower.com setp(pbase(), epptr()); 4912391Sjason@lowepower.com } 5012391Sjason@lowepower.com return 0; 5112391Sjason@lowepower.com } 5212391Sjason@lowepower.com 5312391Sjason@lowepower.compublic: 5412391Sjason@lowepower.com pythonbuf(object pyostream) 5512391Sjason@lowepower.com : pywrite(pyostream.attr("write")), 5612391Sjason@lowepower.com pyflush(pyostream.attr("flush")) { 5712391Sjason@lowepower.com setp(d_buffer, d_buffer + sizeof(d_buffer) - 1); 5812391Sjason@lowepower.com } 5912391Sjason@lowepower.com 6012391Sjason@lowepower.com /// Sync before destroy 6112391Sjason@lowepower.com ~pythonbuf() { 6212391Sjason@lowepower.com sync(); 6312391Sjason@lowepower.com } 6412391Sjason@lowepower.com}; 6512391Sjason@lowepower.com 6612391Sjason@lowepower.comNAMESPACE_END(detail) 6712391Sjason@lowepower.com 6812391Sjason@lowepower.com 6912391Sjason@lowepower.com/** \rst 7012391Sjason@lowepower.com This a move-only guard that redirects output. 7112391Sjason@lowepower.com 7212391Sjason@lowepower.com .. code-block:: cpp 7312391Sjason@lowepower.com 7412391Sjason@lowepower.com #include <pybind11/iostream.h> 7512391Sjason@lowepower.com 7612391Sjason@lowepower.com ... 7712391Sjason@lowepower.com 7812391Sjason@lowepower.com { 7912391Sjason@lowepower.com py::scoped_ostream_redirect output; 8012391Sjason@lowepower.com std::cout << "Hello, World!"; // Python stdout 8112391Sjason@lowepower.com } // <-- return std::cout to normal 8212391Sjason@lowepower.com 8312391Sjason@lowepower.com You can explicitly pass the c++ stream and the python object, 8412391Sjason@lowepower.com for example to guard stderr instead. 8512391Sjason@lowepower.com 8612391Sjason@lowepower.com .. code-block:: cpp 8712391Sjason@lowepower.com 8812391Sjason@lowepower.com { 8912391Sjason@lowepower.com py::scoped_ostream_redirect output{std::cerr, py::module::import("sys").attr("stderr")}; 9012391Sjason@lowepower.com std::cerr << "Hello, World!"; 9112391Sjason@lowepower.com } 9212391Sjason@lowepower.com \endrst */ 9312391Sjason@lowepower.comclass scoped_ostream_redirect { 9412391Sjason@lowepower.comprotected: 9512391Sjason@lowepower.com std::streambuf *old; 9612391Sjason@lowepower.com std::ostream &costream; 9712391Sjason@lowepower.com detail::pythonbuf buffer; 9812391Sjason@lowepower.com 9912391Sjason@lowepower.compublic: 10012391Sjason@lowepower.com scoped_ostream_redirect( 10112391Sjason@lowepower.com std::ostream &costream = std::cout, 10212391Sjason@lowepower.com object pyostream = module::import("sys").attr("stdout")) 10312391Sjason@lowepower.com : costream(costream), buffer(pyostream) { 10412391Sjason@lowepower.com old = costream.rdbuf(&buffer); 10512391Sjason@lowepower.com } 10612391Sjason@lowepower.com 10712391Sjason@lowepower.com ~scoped_ostream_redirect() { 10812391Sjason@lowepower.com costream.rdbuf(old); 10912391Sjason@lowepower.com } 11012391Sjason@lowepower.com 11112391Sjason@lowepower.com scoped_ostream_redirect(const scoped_ostream_redirect &) = delete; 11212391Sjason@lowepower.com scoped_ostream_redirect(scoped_ostream_redirect &&other) = default; 11312391Sjason@lowepower.com scoped_ostream_redirect &operator=(const scoped_ostream_redirect &) = delete; 11412391Sjason@lowepower.com scoped_ostream_redirect &operator=(scoped_ostream_redirect &&) = delete; 11512391Sjason@lowepower.com}; 11612391Sjason@lowepower.com 11712391Sjason@lowepower.com 11812391Sjason@lowepower.com/** \rst 11912391Sjason@lowepower.com Like `scoped_ostream_redirect`, but redirects cerr by default. This class 12012391Sjason@lowepower.com is provided primary to make ``py::call_guard`` easier to make. 12112391Sjason@lowepower.com 12212391Sjason@lowepower.com .. code-block:: cpp 12312391Sjason@lowepower.com 12412391Sjason@lowepower.com m.def("noisy_func", &noisy_func, 12512391Sjason@lowepower.com py::call_guard<scoped_ostream_redirect, 12612391Sjason@lowepower.com scoped_estream_redirect>()); 12712391Sjason@lowepower.com 12812391Sjason@lowepower.com\endrst */ 12912391Sjason@lowepower.comclass scoped_estream_redirect : public scoped_ostream_redirect { 13012391Sjason@lowepower.compublic: 13112391Sjason@lowepower.com scoped_estream_redirect( 13212391Sjason@lowepower.com std::ostream &costream = std::cerr, 13312391Sjason@lowepower.com object pyostream = module::import("sys").attr("stderr")) 13412391Sjason@lowepower.com : scoped_ostream_redirect(costream,pyostream) {} 13512391Sjason@lowepower.com}; 13612391Sjason@lowepower.com 13712391Sjason@lowepower.com 13812391Sjason@lowepower.comNAMESPACE_BEGIN(detail) 13912391Sjason@lowepower.com 14012391Sjason@lowepower.com// Class to redirect output as a context manager. C++ backend. 14112391Sjason@lowepower.comclass OstreamRedirect { 14212391Sjason@lowepower.com bool do_stdout_; 14312391Sjason@lowepower.com bool do_stderr_; 14412391Sjason@lowepower.com std::unique_ptr<scoped_ostream_redirect> redirect_stdout; 14512391Sjason@lowepower.com std::unique_ptr<scoped_estream_redirect> redirect_stderr; 14612391Sjason@lowepower.com 14712391Sjason@lowepower.compublic: 14812391Sjason@lowepower.com OstreamRedirect(bool do_stdout = true, bool do_stderr = true) 14912391Sjason@lowepower.com : do_stdout_(do_stdout), do_stderr_(do_stderr) {} 15012391Sjason@lowepower.com 15112391Sjason@lowepower.com void enter() { 15212391Sjason@lowepower.com if (do_stdout_) 15312391Sjason@lowepower.com redirect_stdout.reset(new scoped_ostream_redirect()); 15412391Sjason@lowepower.com if (do_stderr_) 15512391Sjason@lowepower.com redirect_stderr.reset(new scoped_estream_redirect()); 15612391Sjason@lowepower.com } 15712391Sjason@lowepower.com 15812391Sjason@lowepower.com void exit() { 15912391Sjason@lowepower.com redirect_stdout.reset(); 16012391Sjason@lowepower.com redirect_stderr.reset(); 16112391Sjason@lowepower.com } 16212391Sjason@lowepower.com}; 16312391Sjason@lowepower.com 16412391Sjason@lowepower.comNAMESPACE_END(detail) 16512391Sjason@lowepower.com 16612391Sjason@lowepower.com/** \rst 16712391Sjason@lowepower.com This is a helper function to add a C++ redirect context manager to Python 16812391Sjason@lowepower.com instead of using a C++ guard. To use it, add the following to your binding code: 16912391Sjason@lowepower.com 17012391Sjason@lowepower.com .. code-block:: cpp 17112391Sjason@lowepower.com 17212391Sjason@lowepower.com #include <pybind11/iostream.h> 17312391Sjason@lowepower.com 17412391Sjason@lowepower.com ... 17512391Sjason@lowepower.com 17612391Sjason@lowepower.com py::add_ostream_redirect(m, "ostream_redirect"); 17712391Sjason@lowepower.com 17812391Sjason@lowepower.com You now have a Python context manager that redirects your output: 17912391Sjason@lowepower.com 18012391Sjason@lowepower.com .. code-block:: python 18112391Sjason@lowepower.com 18212391Sjason@lowepower.com with m.ostream_redirect(): 18312391Sjason@lowepower.com m.print_to_cout_function() 18412391Sjason@lowepower.com 18512391Sjason@lowepower.com This manager can optionally be told which streams to operate on: 18612391Sjason@lowepower.com 18712391Sjason@lowepower.com .. code-block:: python 18812391Sjason@lowepower.com 18912391Sjason@lowepower.com with m.ostream_redirect(stdout=true, stderr=true): 19012391Sjason@lowepower.com m.noisy_function_with_error_printing() 19112391Sjason@lowepower.com 19212391Sjason@lowepower.com \endrst */ 19312391Sjason@lowepower.cominline class_<detail::OstreamRedirect> add_ostream_redirect(module m, std::string name = "ostream_redirect") { 19412391Sjason@lowepower.com return class_<detail::OstreamRedirect>(m, name.c_str(), module_local()) 19512391Sjason@lowepower.com .def(init<bool,bool>(), arg("stdout")=true, arg("stderr")=true) 19612391Sjason@lowepower.com .def("__enter__", &detail::OstreamRedirect::enter) 19712391Sjason@lowepower.com .def("__exit__", [](detail::OstreamRedirect &self, args) { self.exit(); }); 19812391Sjason@lowepower.com} 19912391Sjason@lowepower.com 20012391Sjason@lowepower.comNAMESPACE_END(PYBIND11_NAMESPACE) 201