Output not visible in Jupyter notebook

[j-faria]

I'm wondering if something like pybind11's scoped_ostream_redirect is available in nanobind?
Currently, the output from a C++ function which prints to std::cout is only visible in the terminal when running inside a Jupyter notebook.

I tried to implement scoped_ostream_redirect as a call guard, borrowing code from here, with some success.
In the minimal example below I managed to see "before" printed (in both the terminal and a notebook) but not the output from foo() or the "after". Moreover, running in the Jupyter notebook ends up killing the kernel before the function returns.

m.def("foo", [](args...) {
    scoped_ostream_redirect output;
    std::cout << "before" << std::endl; 
    foo();
    std::cout << "after" << std::endl; 
});

My reimplementation is below, almost verbatim from pybind11.

// Buffer that writes to Python instead of C++
class pythonbuf : public std::streambuf {
private:
    using traits_type = std::streambuf::traits_type;

    size_t buf_size;
    std::unique_ptr<char[]> d_buffer;
    nb::object pywrite;
    nb::object pyflush;

    int overflow(int c) override {
        if (!traits_type::eq_int_type(c, traits_type::eof())) {
            *pptr() = traits_type::to_char_type(c);
            pbump(1);
        }
        return sync() == 0 ? traits_type::not_eof(c) : traits_type::eof();
    }

    // Computes how many bytes at the end of the buffer are part of an
    // incomplete sequence of UTF-8 bytes.
    // Precondition: pbase() < pptr()
    size_t utf8_remainder() const {
        const auto rbase = std::reverse_iterator<char *>(pbase());
        const auto rpptr = std::reverse_iterator<char *>(pptr());
        auto is_ascii = [](char c) { return (static_cast<unsigned char>(c) & 0x80) == 0x00; };
        auto is_leading = [](char c) { return (static_cast<unsigned char>(c) & 0xC0) == 0xC0; };
        auto is_leading_2b = [](char c) { return static_cast<unsigned char>(c) <= 0xDF; };
        auto is_leading_3b = [](char c) { return static_cast<unsigned char>(c) <= 0xEF; };
        // If the last character is ASCII, there are no incomplete code points
        if (is_ascii(*rpptr)) {
            return 0;
        }
        // Otherwise, work back from the end of the buffer and find the first
        // UTF-8 leading byte
        const auto rpend = rbase - rpptr >= 3 ? rpptr + 3 : rbase;
        const auto leading = std::find_if(rpptr, rpend, is_leading);
        if (leading == rbase) {
            return 0;
        }
        const auto dist = static_cast<size_t>(leading - rpptr);
        size_t remainder = 0;

        if (dist == 0) {
            remainder = 1; // 1-byte code point is impossible
        } else if (dist == 1) {
            remainder = is_leading_2b(*leading) ? 0 : dist + 1;
        } else if (dist == 2) {
            remainder = is_leading_3b(*leading) ? 0 : dist + 1;
        }
        // else if (dist >= 3), at least 4 bytes before encountering an UTF-8
        // leading byte, either no remainder or invalid UTF-8.
        // Invalid UTF-8 will cause an exception later when converting
        // to a Python string, so that's not handled here.
        return remainder;
    }

    // This function must be non-virtual to be called in a destructor.
    int _sync() {
        if (pbase() != pptr()) { // If buffer is not empty
            //nb::gil_scoped_acquire tmp;
            // This subtraction cannot be negative, so dropping the sign.
            auto size = static_cast<size_t>(pptr() - pbase());
            size_t remainder = utf8_remainder();

            if (size > remainder) {
                nb::str line(pbase(), size - remainder);
                pywrite(std::move(line));
                pyflush();
            }

            // Copy the remainder at the end of the buffer to the beginning:
            if (remainder > 0) {
                std::memmove(pbase(), pptr() - remainder, remainder);
            }
            setp(pbase(), epptr());
            pbump(static_cast<int>(remainder));
        }
        return 0;
    }

    int sync() override { return _sync(); }

public:
    explicit pythonbuf(const nb::object &pyostream, size_t buffer_size = 1024)
        : buf_size(buffer_size), d_buffer(new char[buf_size]), pywrite(pyostream.attr("write")),
          pyflush(pyostream.attr("flush")) {
        setp(d_buffer.get(), d_buffer.get() + buf_size - 1);
    }

    pythonbuf(pythonbuf &&) = default;

    /// Sync before destroy
    ~pythonbuf() override { _sync(); }
};

class scoped_ostream_redirect
{
    protected:
        std::streambuf *old;
        std::ostream &costream;
        pythonbuf buffer;

    public:
        explicit scoped_ostream_redirect(
            std::ostream &costream = std::cout,
            const nb::object &pyostream = nb::module_::import_("sys").attr("stdout"))
        noexcept
        : costream(costream), buffer(pyostream)
        {
            old = costream.rdbuf(&buffer);
        }
        ~scoped_ostream_redirect() {
            costream.rdbuf(old);
        }
        scoped_ostream_redirect(const scoped_ostream_redirect &) = delete;
        scoped_ostream_redirect(scoped_ostream_redirect &&other) = default;
        scoped_ostream_redirect &operator=(const scoped_ostream_redirect &) = delete;
        scoped_ostream_redirect &operator=(scoped_ostream_redirect &&) = delete;
};

[wjakob]

Hum, there are two things I don't like about this code -- first, it depends on the algorithm header, which is out of control especially on C++20 and later versions. I don't wish nanobind to have any dependence on that. Second, this just handles C++ streams (cout) but doesn't redirect plain C prints that go through the FILE* interface. I think that would actually be a much better place to intercept output since the C++ API also ultimately goes through that.

[wjakob]

(The higher-level point is: yes, I think that this is a useful feature and I'd be happy to include a variant of it in nanobind conditional on it being a better fit for this codebase)

[j-faria]

Thank you for the quick reply.
I understand your concerns about this implementation and it would indeed be nice to also redirect C prints.

I think the problems I'm having come from scoped_ostream_redirect not being thread safe. My function foo() above uses threads and that seems to cause the issues (not seeing output from foo() nor the "after" and the kernel crashing). Do you have a suggestion for a workaround?

[keithlostracco]

@j-faria did you ever solve outputting a stream coming from another thread? I'm having the same issue in a Jupyter notebook.

[j-faria]

@keithlostracco, unfortunately no, I haven't been able to solve it yet

[joaospinto]

+1 for supporting std::cout redirection in nanobind. Would be super useful.